commit 118f1fc726524c7ce728cc1474b7679fa4168177 Author: maxiao1 Date: Sat Sep 13 17:00:20 2025 +0800 sglangv0.5.2 & support Qwen3-Next-80B-A3B-Instruct diff --git a/.clang-format-ignore b/.clang-format-ignore new file mode 100644 index 000000000..15c76cc45 --- /dev/null +++ b/.clang-format-ignore @@ -0,0 +1 @@ +sgl-kernel/3rdparty/tensorrt_llm/* diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 000000000..3c7b67cac --- /dev/null +++ b/.devcontainer/Dockerfile @@ -0,0 +1,35 @@ +FROM lmsysorg/sglang:dev + +# Create non-root user with specified UID and GID +# NOTE: Replace with your own UID and GID. This is a workaround from https://github.com/microsoft/vscode-remote-release/issues/49#issuecomment-489060908. +ARG HOST_UID=1003 +ARG HOST_GID=1003 +RUN groupadd -g $HOST_GID devuser && \ + useradd -m -u $HOST_UID -g $HOST_GID -s /bin/zsh devuser + +# Give devuser sudo access +RUN apt-get update && apt-get install -y sudo && \ + echo "devuser ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/devuser && \ + rm -rf /var/lib/apt/lists/* && \ + apt-get clean + +# Set up oh-my-zsh for devuser +RUN cp -r /root/.oh-my-zsh /home/devuser/.oh-my-zsh && \ + cp /root/.zshrc /home/devuser/.zshrc && \ + cp /root/.vimrc /home/devuser/.vimrc && \ + cp /root/.tmux.conf /home/devuser/.tmux.conf && \ + sed -i 's|/root/.oh-my-zsh|/home/devuser/.oh-my-zsh|g' /home/devuser/.zshrc && \ + chown -R devuser:devuser /home/devuser/ + +# Set workspace directory and ownership +WORKDIR /sgl-workspace/sglang +RUN chown -R devuser:devuser /sgl-workspace + +# Switch to devuser +USER devuser + +# Install uv +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +# Install rust +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 000000000..338b10fe6 --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,30 @@ +{ + "name": "sglang", + "build": { + "dockerfile": "Dockerfile" + }, + "remoteUser": "devuser", + "customizations": { + "vscode": { + "extensions": [ + // Python development + "ms-python.python", + "charliermarsh.ruff", + // Rust development + "rust-lang.rust-analyzer", + "tamasfe.even-better-toml" + ] + } + }, + "forwardPorts": [], + "runArgs": [ + "--gpus", + "all" + ], + // The two lines below ensures that your local changes in the sglang + // repo is automatically synced to the sglang pip package installed + // in the dev docker container. You can remove / comment out these + // two lines if you prefer to sync code changes manually. + "workspaceMount": "source=${localWorkspaceFolder},target=/sgl-workspace/sglang,type=bind", + "workspaceFolder": "/sgl-workspace/sglang" +} diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 000000000..030a7293d --- /dev/null +++ b/.editorconfig @@ -0,0 +1,25 @@ +# https://editorconfig.org/ + +root = true + +[*] +charset = utf-8 +end_of_line = lf +indent_style = space +indent_size = 4 +trim_trailing_whitespace = true +insert_final_newline = true + +[*.{json,yaml,yml}] +indent_size = 2 + +[*.md] +indent_size = 2 +x-soft-wrap-text = true + +[*.rst] +indent_size = 4 +x-soft-wrap-text = true + +[Makefile] +indent_style = tab diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 000000000..32435d6ed --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,21 @@ +.github @merrymercy @zhyncs +/docker @zhyncs @HaiShaw @ByronHsu +/python/pyproject.toml @merrymercy @zhyncs +/python/sglang/* @merrymercy @Ying1123 @zhyncs @hnyls2002 +/python/sglang/srt/constrained @hnyls2002 +/python/sglang/srt/disaggregation @ByronHsu @hnyls2002 +/python/sglang/srt/disaggregation/mooncake @ShangmingCai +/python/sglang/srt/distributed @yizhang2077 @merrymercy +/python/sglang/srt/entrypoints @ispobock @CatherineSue @slin1237 @merrymercy +/python/sglang/srt/eplb @fzyzcjy +/python/sglang/srt/function_call @CatherineSue +/python/sglang/srt/layers @merrymercy @Ying1123 @zhyncs @ispobock @HaiShaw @ch-wan @BBuf @kushanam @Edwardf0t1 +/python/sglang/srt/lora @Ying1123 @Fridge003 @lifuhuang +/python/sglang/srt/managers @merrymercy @Ying1123 @hnyls2002 @xiezhq-hermann +/python/sglang/srt/mem_cache @merrymercy @Ying1123 @hnyls2002 @xiezhq-hermann +/python/sglang/srt/model_executor @merrymercy @Ying1123 @hnyls2002 @zhyncs @ispobock +/python/sglang/srt/multimodal @mickqian @JustinTong0323 +/python/sglang/srt/speculative @Ying1123 @merrymercy @rkooo567 @kssteven418 +/sgl-kernel @zhyncs @ispobock @HandH1998 @BBuf @yizhang2077 @merrymercy @FlamingoPg @HaiShaw +/sgl-router @slin1237 @ByronHsu +/test/srt/test_modelopt* @Edwardf0t1 diff --git a/.github/ISSUE_TEMPLATE/1-bug-report.yml b/.github/ISSUE_TEMPLATE/1-bug-report.yml new file mode 100644 index 000000000..5f6734867 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/1-bug-report.yml @@ -0,0 +1,38 @@ +name: 🐞 Bug report +description: Create a report to help us reproduce and fix the bug +title: "[Bug] " +labels: ['Bug'] + +body: +- type: checkboxes + attributes: + label: Checklist + options: + - label: 1. I have searched related issues but cannot get the expected help. + - label: 2. The bug has not been fixed in the latest version. + - label: 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback. + - label: 4. If the issue you raised is not a bug but a question, please raise a discussion at https://github.com/sgl-project/sglang/discussions/new/choose Otherwise, it will be closed. + - label: 5. Please use English, otherwise it will be closed. +- type: textarea + attributes: + label: Describe the bug + description: A clear and concise description of what the bug is. + validations: + required: true +- type: textarea + attributes: + label: Reproduction + description: | + What command or script did you run? Which **model** are you using? + placeholder: | + A placeholder for the command. + validations: + required: true +- type: textarea + attributes: + label: Environment + description: | + Please provide necessary environment information here with `python3 -m sglang.check_env`. Otherwise the issue will be closed. + placeholder: Environment here. + validations: + required: true diff --git a/.github/ISSUE_TEMPLATE/2-feature-request.yml b/.github/ISSUE_TEMPLATE/2-feature-request.yml new file mode 100644 index 000000000..31bc4a127 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/2-feature-request.yml @@ -0,0 +1,23 @@ +name: 🚀 Feature request +description: Suggest an idea for this project +title: "[Feature] " + +body: +- type: checkboxes + attributes: + label: Checklist + options: + - label: 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sglang/discussions/new/choose Otherwise, it will be closed. + - label: 2. Please use English, otherwise it will be closed. +- type: textarea + attributes: + label: Motivation + description: | + A clear and concise description of the motivation of the feature. + validations: + required: true +- type: textarea + attributes: + label: Related resources + description: | + If there is an official code release or third-party implementations, please also provide the information here, which would be very helpful. diff --git a/.github/REVIEWERS.md b/.github/REVIEWERS.md new file mode 100644 index 000000000..ac9ce6102 --- /dev/null +++ b/.github/REVIEWERS.md @@ -0,0 +1,53 @@ +# Area Reviewer + +Here are some reviewers for common areas. You can ping them to review your code if you touch related parts. + +## Hardware platforms +- general @Alcanderian +- AMD GPU @HaiShaw +- Blackwell GPU @kushanam @trevor-m @zhyncs +- CPU @mingfeima + +## Kernel +- general @zhyncs @ispobock @HandH1998 @BBuf @yizhang2077 @HaiShaw +- triton attention backend @ispobock +- aiter attention backend @HaiShaw @kkHuang-amd @valarLip +- flash attention backend @hebiao064 +- flashinfer attention backend @Fridge003 +- moe kernel @BBuf @fzyzcjy @ch-wan @Alcanderian + +## Scheduler and memory pool +- general @merrymercy @Ying1123 @hnyls2002 @xiezhq-hermann +- constrained decoding @hnyls2002 +- hierarchical cache @xiezhq-hermann @DarkSharpness +- lora @Fridge003 @Ying1123 @lifuhuang +- speculative decoding @merrymercy @Ying1123 @kssteven418 @Qiaolin-Yu +- sliding window attention @hanming-lu + +## Parallelism +- expert parallelism @fzyzcjy @ch-wan +- data parallelism attention @ch-wan +- pipeline parallelism @Ying1123 +- tensor parallelism @merrymercy + +## PD disaggregation +- general @ByronHsu @ShangmingCai @hnyls2002 +- Mooncake backend @ShangmingCai + +## Build and release +- general @zhyncs @merrymercy + +## API Server +- general @CatherineSue @slin1237 @ispobock +- function calling and reasoning parsing @CatherineSue +- OpenAI API @CatherineSue @slin1237 + +## SGL-Router +- general @slin1237 @ByronHsu + +## Model +- multimodal models @mickqian @JustinTong0323 +- other new models @zhaochenyang20 + +## Reinforcment learning +- general @zhaochenyang20 @hebiao064 @fzyzcjy @zhuzilin diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 000000000..ab51d4bf5 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,24 @@ + + +## Motivation + + + +## Modifications + + + +## Accuracy Tests + + + +## Benchmarking and Profiling + + + +## Checklist + +- [ ] Format your code according to the [Format code with pre-commit](https://docs.sglang.ai/developer_guide/contribution_guide.html#format-code-with-pre-commit). +- [ ] Add unit tests according to the [Run and add unit tests](https://docs.sglang.ai/developer_guide/contribution_guide.html#run-and-add-unit-tests). +- [ ] Update documentation according to [Write documentations](https://docs.sglang.ai/developer_guide/contribution_guide.html#write-documentations). +- [ ] Provide accuracy and speed benchmark results according to [Test the accuracy](https://docs.sglang.ai/developer_guide/contribution_guide.html#test-the-accuracy) and [Benchmark the speed](https://docs.sglang.ai/developer_guide/contribution_guide.html#benchmark-the-speed). diff --git a/.github/workflows/cancel-all-pending-pr-test-runs.yml b/.github/workflows/cancel-all-pending-pr-test-runs.yml new file mode 100644 index 000000000..6217542e7 --- /dev/null +++ b/.github/workflows/cancel-all-pending-pr-test-runs.yml @@ -0,0 +1,45 @@ +name: Cancel All Pending PR Test Runs + +on: + workflow_dispatch: + inputs: + workflows: + description: 'Space-separated list of workflow filenames to cancel' + required: true + type: string + default: 'pr-test.yml pr-test-xeon.yml' + +permissions: + actions: write # Needed to cancel runs + contents: read # Needed to read repo info + +jobs: + cancel-pending: + runs-on: ubuntu-latest + steps: + - name: Install GitHub CLI + run: sudo apt-get install -y gh jq + + - name: Cancel all pending/waiting runs for specified workflows + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ github.repository }} + run: | + # Read the space-separated string from the input into a bash array + WORKFLOW_FILES=(${{ github.event.inputs.workflows }}) + + echo "Targeting ${#WORKFLOW_FILES[@]} workflow(s): ${{ github.event.inputs.workflows }}" + + for workflow_file in "${WORKFLOW_FILES[@]}"; do + echo "--- Checking workflow: $workflow_file ---" + gh run list \ + --repo "$REPO" \ + --workflow "$workflow_file" \ + --json databaseId,status \ + --limit 1000 \ + | jq -r '.[] | select(.status=="queued" or .status=="in_progress") | .databaseId' \ + | while read run_id; do + echo "Cancelling run ID: $run_id for workflow: $workflow_file" + gh run cancel "$run_id" --repo "$REPO" + done + done diff --git a/.github/workflows/cancel-pr-workflow-on-merge.yml b/.github/workflows/cancel-pr-workflow-on-merge.yml new file mode 100644 index 000000000..535884ba6 --- /dev/null +++ b/.github/workflows/cancel-pr-workflow-on-merge.yml @@ -0,0 +1,22 @@ +name: Cancel PR Workflows on Merge + +on: + pull_request_target: + types: + - closed + +permissions: + actions: write + +jobs: + cancel: + if: github.event.pull_request.merged == true + runs-on: ubuntu-latest + steps: + - name: Cancel Previous Runs + uses: styfle/cancel-workflow-action@0.12.1 + with: + workflow_id: all + access_token: ${{ secrets.GITHUB_TOKEN }} + ignore_sha: true + pr_number: ${{ github.event.pull_request.number }} diff --git a/.github/workflows/close-inactive-issues.yml b/.github/workflows/close-inactive-issues.yml new file mode 100644 index 000000000..048e6c443 --- /dev/null +++ b/.github/workflows/close-inactive-issues.yml @@ -0,0 +1,96 @@ +name: Close Inactive Issues + +on: + schedule: + - cron: '0 0 * * *' + workflow_dispatch: + +permissions: + issues: write + contents: read + +jobs: + close-inactive-issues: + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-latest + steps: + - name: Check and close inactive issues + uses: actions/github-script@v6 + with: + github-token: ${{secrets.GITHUB_TOKEN}} + script: | + const sixtyDaysAgo = new Date(Date.now() - 60 * 24 * 60 * 60 * 1000); + + const [owner, repo] = process.env.GITHUB_REPOSITORY.split('/'); + console.log(`Owner: ${owner}, Repo: ${repo}`); + + async function fetchIssues(page = 1) { + console.log(`Fetching issues for ${owner}/${repo}, page ${page}`); + return await github.rest.issues.listForRepo({ + owner, + repo, + state: 'open', + sort: 'updated', + direction: 'asc', + per_page: 100, + page: page + }); + } + + async function processIssues() { + console.log('Starting to process issues'); + console.log(`Repository: ${owner}/${repo}`); + + let page = 1; + let hasMoreIssues = true; + while (hasMoreIssues) { + try { + const issues = await fetchIssues(page); + console.log(`Fetched ${issues.data.length} issues on page ${page}`); + + if (issues.data.length === 0) { + hasMoreIssues = false; + break; + } + + for (const issue of issues.data) { + // Skip if the issue has 'good first issue' label + if (issue.labels.some(label => label.name === 'good first issue')) { + console.log(`Skipping issue #${issue.number} as it's marked as 'good first issue'`); + continue; + } + if (new Date(issue.updated_at) < sixtyDaysAgo) { + try { + await github.rest.issues.update({ + owner, + repo, + issue_number: issue.number, + state: 'closed', + labels: [...issue.labels.map(l => l.name), 'inactive'] + }); + await github.rest.issues.createComment({ + owner, + repo, + issue_number: issue.number, + body: 'This issue has been automatically closed due to inactivity. Please feel free to reopen it if needed.' + }); + console.log(`Closed issue #${issue.number} due to inactivity.`); + } catch (error) { + console.error(`Failed to close issue #${issue.number}: ${error.message}`); + } + } else { + console.log(`Issue #${issue.number} is still active. Stopping processing.`); + hasMoreIssues = false; + break; + } + } + page += 1; + } catch (error) { + console.error(`Error fetching issues on page ${page}: ${error.message}`); + hasMoreIssues = false; + } + } + console.log('Finished processing issues'); + } + + await processIssues(); diff --git a/.github/workflows/execute-notebook.yml b/.github/workflows/execute-notebook.yml new file mode 100644 index 000000000..7298d80ec --- /dev/null +++ b/.github/workflows/execute-notebook.yml @@ -0,0 +1,60 @@ +name: Execute Notebooks + +on: + pull_request: + branches: [ main ] + paths: + - "python/sglang/**" + - "docs/**" + workflow_dispatch: + + +concurrency: + group: execute-notebook-${{ github.ref }} + cancel-in-progress: true + + +jobs: + run-all-notebooks: + runs-on: 1-gpu-runner + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + pip install -r docs/requirements.txt + apt-get update && apt-get install -y pandoc parallel retry + ln -sf "$(which python3)" /usr/bin/python + + - name: Setup Jupyter Kernel + run: | + python -m ipykernel install --user --name python3 --display-name "Python 3" + + - name: Execute notebooks + timeout-minutes: 40 + run: | + cd docs + make clean + make compile + + + notebook-finish: + needs: [ + run-all-notebooks + ] + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + results=(${{ join(needs.*.result, ' ') }}) + for result in "${results[@]}"; do + if [ "$result" = "failure" ] || [ "$result" = "cancelled" ]; then + echo "Job failed with result: $result" + exit 1 + fi + done + echo "All jobs completed successfully" + exit 0 diff --git a/.github/workflows/experiment-runner.yml b/.github/workflows/experiment-runner.yml new file mode 100644 index 000000000..487ed9ba3 --- /dev/null +++ b/.github/workflows/experiment-runner.yml @@ -0,0 +1,30 @@ +name: Experiment Runner + +on: + workflow_dispatch: + inputs: + script: + description: "Experiment Runner Script" + default: "configs/sharegpt_config.yaml" + +concurrency: + group: experiment-runner-${{ github.ref }} + cancel-in-progress: true + +jobs: + experiment-runner-1-gpu: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + runs-on: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Test experiment runner + timeout-minutes: 120 + run: | + cd test/srt + python3 experiment_runner.py --config ${{ inputs.script }} diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 000000000..3a281299a --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,22 @@ +name: Lint + +on: [pull_request] + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + + - name: Install pre-commit hook + run: | + python -m pip install pre-commit + pre-commit install + + - name: Linting + run: pre-commit run --all-files --show-diff-on-failure diff --git a/.github/workflows/nightly-test-amd.yml b/.github/workflows/nightly-test-amd.yml new file mode 100644 index 000000000..096e876de --- /dev/null +++ b/.github/workflows/nightly-test-amd.yml @@ -0,0 +1,41 @@ +name: Nightly Test (AMD) + +on: + schedule: + - cron: '0 0 * * *' + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +concurrency: + group: nightly-test-amd-${{ github.ref }} + cancel-in-progress: true + +jobs: + nightly-test: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + strategy: + matrix: + runner: [linux-mi300-gpu-2, linux-mi325-gpu-2-nightly] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Setup docker + run: | + touch github_summary.md + bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Nightly Test + run: | + bash scripts/ci/amd_ci_exec.sh -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" python3 run_suite.py --suite nightly-amd --timeout-per-file 7200 + echo "$(> $GITHUB_STEP_SUMMARY diff --git a/.github/workflows/nightly-test.yml b/.github/workflows/nightly-test.yml new file mode 100644 index 000000000..a32c1dbea --- /dev/null +++ b/.github/workflows/nightly-test.yml @@ -0,0 +1,33 @@ +name: Nightly Test + +on: + schedule: + - cron: '0 0 * * *' + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +concurrency: + group: nightly-test-${{ github.ref }} + cancel-in-progress: true + +jobs: + nightly-test: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + runs-on: 2-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 120 + run: | + cd test/srt + python3 run_suite.py --suite nightly --timeout-per-file 3600 diff --git a/.github/workflows/open-pr-copy-from-oss.yml b/.github/workflows/open-pr-copy-from-oss.yml new file mode 100644 index 000000000..05af6ea44 --- /dev/null +++ b/.github/workflows/open-pr-copy-from-oss.yml @@ -0,0 +1,28 @@ +name: Open A PR to Copy Code From OSS + +on: + workflow_dispatch: + # schedule: + # - cron: '0 10 * * *' + +permissions: + contents: write + +jobs: + copy: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + ref: 'main' + + - name: Install GitHub CLI (if not present) + run: | + bash scripts/code_sync/install_github_cli.sh + + - name: Copy from OSS code + env: + GH_TOKEN: ${{ secrets.PAT_FOR_CODE_SYNC_FROM_LIANMIN }} + run: | + python3 scripts/code_sync/copy_from_oss.py diff --git a/.github/workflows/open-pr-copy-to-oss.yml b/.github/workflows/open-pr-copy-to-oss.yml new file mode 100644 index 000000000..b3bb6aae4 --- /dev/null +++ b/.github/workflows/open-pr-copy-to-oss.yml @@ -0,0 +1,31 @@ +name: Open A PR to Copy Diff To OSS + +on: + workflow_dispatch: + inputs: + commit_sha: + description: 'The commit SHA to copy. Defaults to LAST to copy the latest commit.' + required: false + default: 'LAST' + +permissions: + contents: write + +jobs: + copy: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Install GitHub CLI (if not present) + run: | + bash scripts/code_sync/install_github_cli.sh + + - name: Copy to OSS code + env: + GH_TOKEN: ${{ secrets.PAT_FOR_CODE_SYNC_FROM_LIANMIN }} + run: | + python3 scripts/code_sync/copy_to_oss.py --commit ${{ github.event.inputs.commit_sha }} diff --git a/.github/workflows/pr-benchmark-rust.yml b/.github/workflows/pr-benchmark-rust.yml new file mode 100644 index 000000000..e039cba23 --- /dev/null +++ b/.github/workflows/pr-benchmark-rust.yml @@ -0,0 +1,306 @@ +name: PR Benchmark (Rust Router) + +on: + push: + branches: [ main ] + paths: + - "sgl-router/**" + pull_request: + branches: [ main ] + paths: + - "sgl-router/**" + types: [opened, synchronize, reopened, labeled] + workflow_dispatch: + +concurrency: + group: pr-benchmark-rust-${{ github.ref }} + cancel-in-progress: true +permissions: + contents: read + pull-requests: write + issues: write +jobs: + # Quick check job that always runs on PRs + benchmark-compile-check: + name: Benchmark Compilation Check + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.3 + continue-on-error: true + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + # Share cache across all benchmark jobs + shared-key: "rust-cache" + # Save cache even on failure + save-if: true + + - name: Check benchmarks compile + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + # Try to use sccache, but disable if it fails + if command -v sccache &> /dev/null; then + echo "Testing sccache availability..." + # Try to start sccache and check if it works + export RUSTC_WRAPPER=sccache + export SCCACHE_GHA_ENABLED="true" + if sccache --start-server 2>/dev/null && sccache --show-stats 2>/dev/null; then + echo "sccache is working, using it for compilation" + else + echo "sccache failed to start, falling back to regular cargo" + unset RUSTC_WRAPPER + unset SCCACHE_GHA_ENABLED + fi + else + echo "sccache not available, using regular cargo" + fi + cargo check --benches + + # Full benchmark jobs that only run with label or on main branch + benchmark-request-processing: + name: Request Processing Benchmark + if: | + github.repository == 'sgl-project/sglang' && + (github.event_name == 'push' || + github.event_name == 'workflow_dispatch' || + contains(github.event.pull_request.labels.*.name, 'benchmark')) + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + # Fetch enough history for baseline comparison + fetch-depth: 100 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.3 + continue-on-error: true + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + # Share cache across all benchmark jobs + shared-key: "rust-cache" + # Save cache even on failure + save-if: true + + - name: Run request processing benchmark + timeout-minutes: 30 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + # Try to use sccache, but disable if it fails + if command -v sccache &> /dev/null; then + echo "Testing sccache availability..." + # Try to start sccache and check if it works + export RUSTC_WRAPPER=sccache + export SCCACHE_GHA_ENABLED="true" + if sccache --start-server 2>/dev/null && sccache --show-stats 2>/dev/null; then + echo "sccache is working, using it for compilation" + else + echo "sccache failed to start, falling back to regular cargo" + unset RUSTC_WRAPPER + unset SCCACHE_GHA_ENABLED + fi + else + echo "sccache not available, using regular cargo" + fi + # Run only the summary benchmark for quick validation in PRs + cargo bench --bench request_processing -- benchmark_summary --exact + + - name: Upload benchmark results + if: always() + uses: actions/upload-artifact@v4 + with: + name: request-processing-results-${{ github.sha }} + path: | + sgl-router/target/criterion/benchmark_summary/ + retention-days: 30 + + benchmark-tokenizer: + name: Tokenizer Benchmark + if: | + github.repository == 'sgl-project/sglang' && + (github.event_name == 'push' || + github.event_name == 'workflow_dispatch' || + contains(github.event.pull_request.labels.*.name, 'benchmark')) + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 100 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.3 + continue-on-error: true + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + # Share cache across all benchmark jobs + shared-key: "rust-cache" + # Save cache even on failure + save-if: true + + - name: Run tokenizer benchmark + timeout-minutes: 30 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + # Try to use sccache, but disable if it fails + if command -v sccache &> /dev/null; then + echo "Testing sccache availability..." + # Try to start sccache and check if it works + export RUSTC_WRAPPER=sccache + export SCCACHE_GHA_ENABLED="true" + if sccache --start-server 2>/dev/null && sccache --show-stats 2>/dev/null; then + echo "sccache is working, using it for compilation" + else + echo "sccache failed to start, falling back to regular cargo" + unset RUSTC_WRAPPER + unset SCCACHE_GHA_ENABLED + fi + else + echo "sccache not available, using regular cargo" + fi + cargo bench --bench tokenizer_benchmark + + - name: Upload benchmark results + if: always() + uses: actions/upload-artifact@v4 + with: + name: tokenizer-results-${{ github.sha }} + path: | + sgl-router/target/criterion/tokenizer*/ + retention-days: 30 + + benchmark-tool-parser: + name: Tool Parser Benchmark + if: | + github.repository == 'sgl-project/sglang' && + (github.event_name == 'push' || + github.event_name == 'workflow_dispatch' || + contains(github.event.pull_request.labels.*.name, 'benchmark')) + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 100 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.3 + continue-on-error: true + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + # Share cache across all benchmark jobs + shared-key: "rust-cache" + # Save cache even on failure + save-if: true + + - name: Run tool parser benchmark + timeout-minutes: 30 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + # Try to use sccache, but disable if it fails + if command -v sccache &> /dev/null; then + echo "Testing sccache availability..." + # Try to start sccache and check if it works + export RUSTC_WRAPPER=sccache + export SCCACHE_GHA_ENABLED="true" + if sccache --start-server 2>/dev/null && sccache --show-stats 2>/dev/null; then + echo "sccache is working, using it for compilation" + else + echo "sccache failed to start, falling back to regular cargo" + unset RUSTC_WRAPPER + unset SCCACHE_GHA_ENABLED + fi + else + echo "sccache not available, using regular cargo" + fi + cargo bench --bench tool_parser_benchmark + + - name: Upload benchmark results + if: always() + uses: actions/upload-artifact@v4 + with: + name: tool-parser-results-${{ github.sha }} + path: | + sgl-router/target/criterion/tool_parser*/ + retention-days: 30 + + benchmark-summary: + name: Benchmark Summary + needs: [benchmark-request-processing, benchmark-tokenizer, benchmark-tool-parser] + if: always() && (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') + runs-on: ubuntu-latest + steps: + - name: Download all benchmark results + uses: actions/download-artifact@v4 + with: + pattern: '*-results-${{ github.sha }}' + path: benchmark-results + + - name: Generate summary + run: | + echo "## Benchmark Results Summary" > summary.md + echo "" >> summary.md + echo "### Request Processing" >> summary.md + if [ -d "benchmark-results/request-processing-results-${{ github.sha }}" ]; then + echo "✅ Completed" >> summary.md + else + echo "❌ Failed or skipped" >> summary.md + fi + echo "" >> summary.md + echo "### Tokenizer" >> summary.md + if [ -d "benchmark-results/tokenizer-results-${{ github.sha }}" ]; then + echo "✅ Completed" >> summary.md + else + echo "❌ Failed or skipped" >> summary.md + fi + echo "" >> summary.md + echo "### Tool Parser" >> summary.md + if [ -d "benchmark-results/tool-parser-results-${{ github.sha }}" ]; then + echo "✅ Completed" >> summary.md + else + echo "❌ Failed or skipped" >> summary.md + fi + cat summary.md + + - name: Upload summary + uses: actions/upload-artifact@v4 + with: + name: benchmark-summary-${{ github.sha }} + path: summary.md + retention-days: 30 diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml new file mode 100644 index 000000000..2c7e2c652 --- /dev/null +++ b/.github/workflows/pr-test-amd.yml @@ -0,0 +1,377 @@ +name: PR Test (AMD) + +on: + push: + branches: [ main ] + paths: + - "python/**" + - "scripts/ci/**" + - "test/**" + - "sgl-kernel/**" + - ".github/workflows/pr-test-amd.yml" + pull_request: + branches: [ main ] + paths: + - "python/**" + - "scripts/ci/**" + - "test/**" + - "sgl-kernel/**" + - ".github/workflows/pr-test-amd.yml" + workflow_dispatch: + +concurrency: + group: pr-test-amd-${{ github.ref }} + cancel-in-progress: true + +jobs: + accuracy-test-1-gpu-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1, linux-mi325-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Evaluate Accuracy + timeout-minutes: 30 + run: | + bash scripts/ci/amd_ci_exec.sh -e SGLANG_USE_AITER=0 python3 test_eval_accuracy_large.py + bash scripts/ci/amd_ci_exec.sh python3 test_eval_fp8_accuracy.py + bash scripts/ci/amd_ci_exec.sh python3 models/test_qwen_models.py + + accuracy-test-2-gpu-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-2, linux-mi325-gpu-2, linux-mi35x-gpu-2] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Evaluate accuracy (TP=2) + timeout-minutes: 60 + run: | + bash scripts/ci/amd_ci_exec.sh -e SGLANG_USE_AITER=0 python3 test_moe_eval_accuracy_large.py + + mla-test-1-gpu-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1, linux-mi325-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: MLA TEST + timeout-minutes: 30 + run: | + bash scripts/ci/amd_ci_exec.sh python3 test_mla.py + + performance-test-1-gpu-part-1-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1, linux-mi325-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Benchmark single latency + timeout-minutes: 20 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_bs1_small + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_bs1_default + + - name: Benchmark online latency + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_online_latency_default + + - name: Benchmark offline throughput + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_default + + - name: Benchmark offline throughput (Non-streaming, small batch size) + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_non_stream_small_batch_size + + performance-test-1-gpu-part-2-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1, linux-mi325-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Benchmark offline throughput (w/o RadixAttention) + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_without_radix_cache + + - name: Benchmark offline throughput (w/ Triton) + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_with_triton_attention_backend + + - name: Benchmark offline throughput (w/ FP8) + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_default_fp8 + + bench-test-2-gpu-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-2, linux-mi325-gpu-2] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Benchmark dummy grok (TP=2) + timeout-minutes: 30 + run: | + bash scripts/ci/amd_ci_exec.sh python3 models/test_dummy_grok_models.py + + - name: Benchmark single latency (TP=2) + timeout-minutes: 25 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_moe_tp2_bs1 + + - name: Benchmark single latency + torch.compile (TP=2) + timeout-minutes: 25 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_torch_compile_tp2_bs1 + + - name: Benchmark offline throughput (TP=2) + timeout-minutes: 25 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_default + + - name: Benchmark offline throughput (w/o RadixAttention) (TP=2) + timeout-minutes: 25 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_without_radix_cache + + unit-test-backend-1-gpu-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1, linux-mi325-gpu-1] + part: [0, 1, 2, 3, 4, 5, 6, 7] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 50 + run: | + bash scripts/ci/amd_ci_exec.sh python3 run_suite.py --suite per-commit-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 8 + + unit-test-backend-1-gpu-amd-mi35x: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi35x-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 50 + run: | + bash scripts/ci/amd_ci_exec.sh python3 run_suite.py --suite per-commit-amd-mi35x + + unit-test-backend-2-gpu-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-2, linux-mi325-gpu-2] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 40 + run: | + bash scripts/ci/amd_ci_exec.sh python3 run_suite.py --suite per-commit-2-gpu-amd + + unit-test-backend-8-gpu-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-8] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 60 + run: | + bash scripts/ci/amd_ci_exec.sh python3 run_suite.py --suite per-commit-8-gpu-amd --timeout-per-file 3600 + + unit-test-sgl-kernel-amd: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1, linux-mi325-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 14 + run: | + docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_align.py + docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_topk_softmax.py + docker exec -w /sglang-checkout/sgl-kernel/tests/speculative ci_sglang python3 -m pytest test_eagle_utils.py + docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_apply_token_bitmask_inplace.py + docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_activation.py + docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_kvcacheio.py + + pr-test-amd-finish: + if: always() + needs: [ + accuracy-test-1-gpu-amd, mla-test-1-gpu-amd, bench-test-2-gpu-amd, + accuracy-test-2-gpu-amd, performance-test-1-gpu-part-1-amd, performance-test-1-gpu-part-2-amd, + unit-test-backend-1-gpu-amd, unit-test-backend-1-gpu-amd-mi35x, unit-test-backend-2-gpu-amd, + unit-test-backend-8-gpu-amd, unit-test-sgl-kernel-amd + ] + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + results=(${{ join(needs.*.result, ' ') }}) + for result in "${results[@]}"; do + if [ "$result" = "failure" ] || [ "$result" = "cancelled" ]; then + echo "Job failed with result: $result" + exit 1 + fi + done + echo "All jobs completed successfully" + exit 0 diff --git a/.github/workflows/pr-test-h20.yml b/.github/workflows/pr-test-h20.yml new file mode 100644 index 000000000..58e335289 --- /dev/null +++ b/.github/workflows/pr-test-h20.yml @@ -0,0 +1,81 @@ +name: PR Test (H20) + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + inputs: + version: + required: true + type: choice + default: 'release' + options: + - 'release' + - 'nightly' + +concurrency: + group: pr-test-h20-${{ github.ref }} + cancel-in-progress: true + +jobs: + check-changes: + runs-on: ubuntu-latest + outputs: + src: ${{ steps.filter.outputs.src }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + with: + filters: | + src: + - "python/sglang/srt/models/deepseek*" + - "python/sglang/srt/layers/moe/**" + - ".github/workflows/pr-test-h20.yml" + - "python/pyproject.toml" + + per-commit-8-gpu-h20: + needs: [check-changes] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 8-gpu-h20 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 20 + + run: | + cd test/srt + python3 run_suite.py --suite per-commit-8-gpu-h20 + + pr-test-h20-finish: + needs: [ + check-changes, + per-commit-8-gpu-h20, + ] + if: needs.check-changes.outputs.src == 'true' + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + results=(${{ join(needs.*.result, ' ') }}) + for result in "${results[@]}"; do + if [ "$result" = "failure" ] || [ "$result" = "cancelled" ]; then + echo "Job failed with result: $result" + exit 1 + fi + done + echo "All jobs completed successfully" + exit 0 diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml new file mode 100644 index 000000000..c0fe381e3 --- /dev/null +++ b/.github/workflows/pr-test-npu.yml @@ -0,0 +1,184 @@ +name: PR Test (Ascend NPU) + +on: + push: + branches: [ main ] + paths: + - "python/**" + - "scripts/ci/**" + - "test/**" + - ".github/workflows/pr-test-npu.yml" + pull_request: + branches: [ main ] + paths: + - "python/**" + - "scripts/ci/**" + - "test/**" + - ".github/workflows/pr-test-npu.yml" + workflow_dispatch: + +concurrency: + group: pr-test-npu-${{ github.ref }} + cancel-in-progress: true + +jobs: + per-commit-1-ascend-npu: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + runs-on: linux-arm64-npu-1 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.2.rc1-910b-ubuntu22.04-py3.11 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host ${CACHING_URL} + + bash scripts/ci/npu_ci_install_dependency.sh + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy download through proxy + curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl + + - name: Run test + timeout-minutes: 60 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + run: | + cd test/srt + python3 run_suite.py --suite per-commit-1-ascend-npu + + per-commit-2-ascend-npu: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + runs-on: linux-arm64-npu-2 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.2.rc1-910b-ubuntu22.04-py3.11 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host ${CACHING_URL} + + bash scripts/ci/npu_ci_install_dependency.sh + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy download through proxy + curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl + + - name: Run test + timeout-minutes: 90 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + run: | + cd test/srt + python3 run_suite.py --suite per-commit-2-ascend-npu + + per-commit-4-ascend-npu: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + runs-on: linux-arm64-npu-4 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.2.rc1-910b-ubuntu22.04-py3.11 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host ${CACHING_URL} + + bash scripts/ci/npu_ci_install_dependency.sh + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy download through proxy + curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl + + - name: Run test + timeout-minutes: 120 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + run: | + cd test/srt + python3 run_suite.py --suite per-commit-4-ascend-npu --timeout-per-file 3600 + + per-commit-16-ascend-a3: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + runs-on: linux-aarch64-a3-16 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.2.rc1-a3-ubuntu22.04-py3.11 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host ${CACHING_URL} + + bash scripts/ci/npu_ci_install_dependency.sh + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy download through proxy + curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl + + - name: Run test + timeout-minutes: 90 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + run: | + cd test/srt + python3 run_suite.py --suite per-commit-16-ascend-a3 --timeout-per-file 5400 + + pr-test-npu-finish: + if: always() + needs: + - per-commit-1-ascend-npu + - per-commit-2-ascend-npu + - per-commit-4-ascend-npu + - per-commit-16-ascend-a3 + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + results=(${{ join(needs.*.result, ' ') }}) + for result in "${results[@]}"; do + if [ "$result" = "failure" ] || [ "$result" = "cancelled" ]; then + echo "Job failed with result: $result" + exit 1 + fi + done + echo "All jobs completed successfully" + exit 0 diff --git a/.github/workflows/pr-test-pd-router.yml b/.github/workflows/pr-test-pd-router.yml new file mode 100644 index 000000000..2a1bde1b4 --- /dev/null +++ b/.github/workflows/pr-test-pd-router.yml @@ -0,0 +1,599 @@ +name: PR Test (PD Router) + +on: + push: + branches: [ main ] + paths: + - 'python/sglang/srt/disaggregation/**' + - 'scripts/ci/ci_start_disaggregation_servers.sh' + - 'sgl-router/**' + pull_request: + branches: [ main ] + paths: + - 'python/sglang/srt/disaggregation/**' + - 'scripts/ci/ci_start_disaggregation_servers.sh' + - 'sgl-router/**' + workflow_dispatch: + +concurrency: + group: test-disaggregation-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + pull-requests: write + issues: write + +jobs: + test-disaggregation: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + runs-on: [h200] + timeout-minutes: 45 + + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 10 + + - name: Setup Python + uses: actions/setup-python@v4 + with: + python-version: '3.12' + + - name: Setup Rust + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Cache Rust dependencies + uses: actions/cache@v4 + with: + path: | + ~/.cargo/bin/ + ~/.cargo/registry/index/ + ~/.cargo/registry/cache/ + ~/.cargo/git/db/ + sgl-router/target/ + key: ${{ runner.os }}-cargo-${{ hashFiles('sgl-router/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo- + + - name: Cache pip dependencies + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('python/pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-pip- + + - name: Validate environment + run: | + echo "=== System Validation ===" + nvidia-smi + echo "GPU count: $(nvidia-smi -L | wc -l)" + if [ $(nvidia-smi -L | wc -l) -lt 8 ]; then + echo "Error: This test requires at least 8 GPUs" + exit 1 + fi + + echo "=== GPU Process Check ===" + # Fail fast if any GPU compute processes are active + if command -v nvidia-smi >/dev/null 2>&1; then + # Try to query compute apps first (preferred and concise) + gpu_procs=$(nvidia-smi --query-compute-apps=pid,process_name,gpu_uuid --format=csv,noheader 2>/dev/null | sed '/^$/d' || true) + + # Fallback to detailed PIDS report if the query returns nothing but there might still be processes + if [ -z "$gpu_procs" ]; then + gpu_procs=$(nvidia-smi -q -d PIDS 2>/dev/null | awk '/Processes/{flag=1;next}/^$/{flag=0}flag' | sed '/^\s*Processes:/d' | sed '/^\s*$/d' || true) + fi + + if [ -n "$gpu_procs" ]; then + echo "Error: Found active GPU processes using the device(s):" + echo "$gpu_procs" + exit 1 + else + echo "No active GPU compute processes detected." + fi + else + echo "Error: nvidia-smi not found; skipping GPU process check." + exit 1 + fi + + echo "=== RDMA Validation ===" + if ! command -v ibv_devices >/dev/null 2>&1; then + echo "Error: InfiniBand tools not found" + exit 1 + fi + + # Check for active IB devices + found_active_device=false + for device in mlx5_{0..11}; do + if ibv_devinfo $device >/dev/null 2>&1; then + state=$(ibv_devinfo $device | grep "state:" | head -1 | awk '{print $2}') + if [[ "$state" == "PORT_ACTIVE" ]]; then + echo "✓ Found active device: $device" + found_active_device=true + break + fi + fi + done + + if [ "$found_active_device" = false ]; then + echo "Error: No active IB devices found" + echo "Available devices:" + ibv_devices || true + exit 1 + fi + + echo "=== Model Validation ===" + if [ ! -d "/raid/models/meta-llama/Llama-3.1-8B-Instruct" ]; then + echo "Error: Model not found" + ls -la /raid/models/ || echo "No models directory" + exit 1 + fi + echo "✓ Model found" + + - name: Install SGLang dependencies + run: | + echo "Installing SGLang with all extras..." + python3 -m pip --no-cache-dir install --upgrade pip + python3 -m pip --no-cache-dir install torch==2.8.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu126 + python3 -m pip --no-cache-dir install -e "python[all]" --break-system-packages + python3 -m pip --no-cache-dir install mooncake-transfer-engine==0.3.5 + python3 -m pip --no-cache-dir install --user --force-reinstall genai-bench==0.0.2 + python3 -m pip --no-cache-dir install sgl-kernel==0.3.9.post2 + + - name: Build and install sgl-router + run: | + source "$HOME/.cargo/env" + echo "Building sgl-router..." + cd sgl-router + cargo build && python3 -m build && pip install --force-reinstall dist/*.whl + + - name: Start disaggregation servers + id: start_servers + run: | + echo "Starting disaggregation servers..." + bash scripts/ci/ci_start_disaggregation_servers.sh & + SERVER_PID=$! + echo "server_pid=$SERVER_PID" >> $GITHUB_OUTPUT + + # Wait for all 8 servers to be healthy (script already does this) + wait_count=0 + while [ $wait_count -lt 30 ]; do + if ps -p $SERVER_PID > /dev/null; then + # Check if the startup script printed success message + sleep 2 + wait_count=$((wait_count + 1)) + else + # Script exited - check if it was successful + wait $SERVER_PID + exit_code=$? + if [ $exit_code -eq 0 ]; then + echo "✓ All disaggregation servers are healthy" + break + else + echo "Error: Server startup failed with code $exit_code" + exit 1 + fi + fi + done + + echo "✓ Servers started (PID: $SERVER_PID)" + + - name: Test all policies sequentially + timeout-minutes: 30 + run: | + POLICIES=("random" "round_robin" "cache_aware" "power_of_two") + BASE_URL="http://127.0.0.9:8000" + + # Free commonly used ports for router and metrics + echo "Freeing ports 29000 (metrics) and 8000 (API), if in use..." + fuser -k -n tcp 29000 2>/dev/null || true + fuser -k -n tcp 8000 2>/dev/null || true + sleep 1 + + for policy in "${POLICIES[@]}"; do + echo "" + echo "==================================================" + echo "Testing policy: $policy" + echo "==================================================" + + # Free ports before starting router + fuser -k -n tcp 29000 2>/dev/null || true + fuser -k -n tcp 8000 2>/dev/null || true + + # Start router with the current policy + echo "Starting router with policy: $policy..." + RUST_BACKTRACE=1 python3 -m sglang_router.launch_router \ + --pd-disaggregation \ + --policy "$policy" \ + --prefill http://127.0.0.1:30001 9001 \ + --prefill http://127.0.0.2:30002 9002 \ + --prefill http://127.0.0.3:30003 9003 \ + --prefill http://127.0.0.4:30004 9004 \ + --decode http://127.0.0.5:30005 \ + --decode http://127.0.0.6:30006 \ + --decode http://127.0.0.7:30007 \ + --decode http://127.0.0.8:30008 \ + --host 127.0.0.9 \ + --port 8000 & + ROUTER_PID=$! + + # Wait for router to become healthy + echo "Waiting for router to become healthy..." + TIMEOUT=60 + ELAPSED=0 + while [ $ELAPSED -lt $TIMEOUT ]; do + if curl --connect-timeout 5 --silent http://127.0.0.9:8000 > /dev/null 2>&1; then + echo "✓ Router is reachable" + break + fi + if ! ps -p $ROUTER_PID > /dev/null; then + echo "Error: Router process died" + exit 1 + fi + sleep 5 + ELAPSED=$((ELAPSED + 5)) + done + + if [ $ELAPSED -ge $TIMEOUT ]; then + echo "Error: Router health check timeout" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + + # Test API functionality + echo "Testing API completions for $policy..." + response=$(curl -s -X POST "$BASE_URL/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer test-token" \ + -d '{ + "model": "/raid/models/meta-llama/Llama-3.1-8B-Instruct", + "messages": [ + {"role": "user", "content": "Write a Python function to calculate fibonacci numbers recursively"} + ], + "stream": false, + "max_tokens": 100 + }') + + if echo "$response" | jq -e '.choices[0].message.content' > /dev/null 2>&1; then + echo "✓ API test passed for $policy" + else + echo "✗ API test failed for $policy: $response" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + + # Test streaming + echo "Testing streaming API for $policy..." + stream_response=$(timeout 30 curl -s -X POST "$BASE_URL/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer test-token" \ + -d '{ + "model": "/raid/models/meta-llama/Llama-3.1-8B-Instruct", + "messages": [ + {"role": "user", "content": "Count from 1 to 5"} + ], + "stream": true, + "max_tokens": 50 + }') + + if echo "$stream_response" | grep -q "data:"; then + echo "✓ Streaming API test passed for $policy" + else + echo "✗ Streaming API test failed for $policy" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + + # Run genai-bench benchmark + echo "Running genai-bench for $policy..." + genai-bench benchmark \ + --api-backend openai \ + --api-base "http://127.0.0.9:8000" \ + --api-key "dummy-token" \ + --api-model-name "/raid/models/meta-llama/Llama-3.1-8B-Instruct" \ + --model-tokenizer /raid/models/meta-llama/Llama-3.1-8B-Instruct \ + --task text-to-text \ + --num-concurrency 64 \ + --traffic-scenario "D(8000,2000)" \ + --max-requests-per-run 640 \ + --max-time-per-run 2 \ + --experiment-folder-name "benchmark_${policy}" \ + --experiment-base-dir "." + + # Find the actual experiment folder + actual_folder=$(find . -maxdepth 1 -name "benchmark_${policy}" -type d | head -1) + + if [ -n "$actual_folder" ]; then + # Extract metrics from the Excel summary or JSON files + summary_file="$actual_folder"/*_summary.xlsx + json_files=$(find "$actual_folder" -name "*.json" | grep -v experiment_metadata) + + echo "Genai-bench results saved in: $actual_folder" + + # Extract mean values and validate performance thresholds + echo "📊 Extracting performance metrics for $policy..." + + # Find JSON files excluding experiment metadata + json_files=$(find "$actual_folder" -name "*.json" | grep -v experiment_metadata) + + if [ -n "$json_files" ]; then + # Extract metrics using jq and validate against loose thresholds + for json_file in $json_files; do + echo "Processing: $(basename "$json_file")" + + # Extract mean values for performance validation + ttft_mean=$(jq -r '.aggregated_metrics.stats.ttft.mean' "$json_file") + e2e_latency_mean=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean' "$json_file") + input_throughput_mean=$(jq -r '.aggregated_metrics.stats.input_throughput.mean' "$json_file") + output_throughput_mean=$(jq -r '.aggregated_metrics.stats.output_throughput.mean' "$json_file") + + echo " TTFT mean: ${ttft_mean}s" + echo " E2E Latency mean: ${e2e_latency_mean}s" + echo " Input Throughput mean: ${input_throughput_mean} tokens/s" + echo " Output Throughput mean: ${output_throughput_mean} tokens/s" + + # Set mean thresholds (allowing for reasonable variance) + # These can be adjusted based on your performance requirements + ttft_threshold=4.7 # Max 4.7 seconds for mean TTFT + e2e_latency_threshold=35.0 # Max 35.0 seconds for mean E2E latency + input_throughput_threshold=12000 # Min 12000 tokens/s for mean input throughput + output_throughput_threshold=68 # Min 68 tokens/s for mean output throughput + + + # Validate mean thresholds + validation_passed=true + + if (( $(echo "$ttft_mean > $ttft_threshold" | bc -l) )); then + echo "❌ TTFT validation failed: $ttft_mean > $ttft_threshold" + validation_passed=false + fi + + if (( $(echo "$e2e_latency_mean > $e2e_latency_threshold" | bc -l) )); then + echo "❌ E2E Latency validation failed: $e2e_latency_mean > $e2e_latency_threshold" + validation_passed=false + fi + + if (( $(echo "$input_throughput_mean < $input_throughput_threshold" | bc -l) )); then + echo "❌ Input Throughput validation failed: $input_throughput_mean < $input_throughput_threshold" + validation_passed=false + fi + + if (( $(echo "$output_throughput_mean < $output_throughput_threshold" | bc -l) )); then + echo "❌ Output Throughput validation failed: $output_throughput_mean < $output_throughput_threshold" + validation_passed=false + fi + + if [ "$validation_passed" = true ]; then + echo "✅ Performance validation passed for $policy" + else + echo "❌ Performance validation failed for $policy" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + done + + echo "✓ Genai-bench completed successfully for $policy" + echo "📊 Detailed metrics and plots available in: $actual_folder" + else + echo "✗ Benchmark failed for $policy: No JSON results found" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + else + echo "✗ Benchmark failed for $policy: Experiment folder not found" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + + # Stop router before testing next policy + echo "Stopping router for $policy..." + # First try graceful shutdown + kill $ROUTER_PID 2>/dev/null || true + + # Wait up to 5 seconds for graceful shutdown + for i in {1..5}; do + if ! ps -p $ROUTER_PID > /dev/null 2>&1; then + echo "Router stopped gracefully" + break + fi + sleep 1 + done + + # Force kill if still running + if ps -p $ROUTER_PID > /dev/null 2>&1; then + echo "Force killing router..." + kill -9 $ROUTER_PID 2>/dev/null || true + fi + + # Short delay to ensure port is released + sleep 2 + + echo "✓ Completed testing for $policy" + done + + echo "" + echo "✅ All policies tested successfully!" + + + - name: Upload benchmark results + if: success() + uses: actions/upload-artifact@v4 + with: + name: genai-bench-results-all-policies + path: benchmark_**/ + + - name: Cleanup servers + if: always() + run: | + if [ -n "${{ steps.start_servers.outputs.server_pid }}" ]; then + pkill -P ${{ steps.start_servers.outputs.server_pid }} || true + kill ${{ steps.start_servers.outputs.server_pid }} || true + fi + pkill -f "sglang.launch_server" || true + sleep 5 + remaining=$(ps aux | grep -c "sglang.launch_server" || echo "0") + echo "Cleanup completed. Remaining processes: $remaining" + + summarize-benchmarks: + needs: test-disaggregation + runs-on: ubuntu-latest + if: success() + + steps: + - name: Install jq + run: sudo apt-get update && sudo apt-get install -y jq bc + + - name: Download benchmark results + uses: actions/download-artifact@v4 + with: + name: genai-bench-results-all-policies + + - name: List downloaded contents + run: | + echo "Contents after download:" + ls -la + find . -name "benchmark_*" -type d + echo "JSON files found:" + find . -name "*.json" | head -10 + + - name: Create benchmark summary + run: | + echo "=== DEBUG: Creating benchmark summary ===" + echo "Available benchmark directories:" + find . -name "benchmark_*" -type d + echo "==========================================" + + echo "## PD Router Genai-Bench Results Summary" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "🚀 **Benchmarked with genai-bench for comprehensive LLM serving performance evaluation**" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "| Policy | Status | TTFT (s) | E2E Latency (s) | Input Throughput (tok/s) | Output Throughput (tok/s) |" >> $GITHUB_STEP_SUMMARY + echo "|--------|--------|----------|-----------------|--------------------------|---------------------------|" >> $GITHUB_STEP_SUMMARY + + # First, complete the table with all policies + for policy in random round_robin cache_aware power_of_two; do + # Find genai-bench result folders for this policy (handle zip extraction structure) + result_folder=$(find . -maxdepth 2 -name "benchmark_${policy}" -type d | head -1) + if [ -z "$result_folder" ]; then + # Try alternative patterns in case of different extraction structure + result_folder=$(find . -maxdepth 3 -path "*benchmark_${policy}*" -type d | head -1) + fi + + echo "DEBUG: Policy ${policy} -> Found folder: ${result_folder:-'NOT FOUND'}" + + if [ -n "$result_folder" ] && [ -d "$result_folder" ]; then + # Find JSON file with metrics + json_file=$(find "$result_folder" -name "*.json" -not -name "experiment_metadata.json" | head -1) + + if [ -n "$json_file" ] && [ -f "$json_file" ]; then + # Extract performance metrics + ttft_mean=$(jq -r '.aggregated_metrics.stats.ttft.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + e2e_latency_mean=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + input_throughput_mean=$(jq -r '.aggregated_metrics.stats.input_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + output_throughput_mean=$(jq -r '.aggregated_metrics.stats.output_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + + # Format numbers for display (2 decimal places) + if [ "$ttft_mean" != "N/A" ] && [ "$ttft_mean" != "null" ]; then + ttft_display=$(printf "%.2f" "$ttft_mean" 2>/dev/null || echo "$ttft_mean") + else + ttft_display="N/A" + fi + + if [ "$e2e_latency_mean" != "N/A" ] && [ "$e2e_latency_mean" != "null" ]; then + e2e_display=$(printf "%.2f" "$e2e_latency_mean" 2>/dev/null || echo "$e2e_latency_mean") + else + e2e_display="N/A" + fi + + if [ "$input_throughput_mean" != "N/A" ] && [ "$input_throughput_mean" != "null" ]; then + input_display=$(printf "%.0f" "$input_throughput_mean" 2>/dev/null || echo "$input_throughput_mean") + else + input_display="N/A" + fi + + if [ "$output_throughput_mean" != "N/A" ] && [ "$output_throughput_mean" != "null" ]; then + output_display=$(printf "%.0f" "$output_throughput_mean" 2>/dev/null || echo "$output_throughput_mean") + else + output_display="N/A" + fi + + echo "| ${policy} | ✅ Success | $ttft_display | $e2e_display | $input_display | $output_display |" >> $GITHUB_STEP_SUMMARY + else + echo "| ${policy} | ❌ No Data | N/A | N/A | N/A | N/A |" >> $GITHUB_STEP_SUMMARY + fi + else + echo "| ${policy} | ❌ Failed | N/A | N/A | N/A | N/A |" >> $GITHUB_STEP_SUMMARY + fi + done + + # Add performance validation summary + echo "" >> $GITHUB_STEP_SUMMARY + echo "## 📊 Performance Validation" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "**Thresholds:** TTFT ≤ 2.0s | E2E Latency ≤ 8.0s | Input Throughput ≥ 10,000 tok/s | Output Throughput ≥ 100 tok/s" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + + validation_summary="" + for policy in random round_robin cache_aware power_of_two; do + # Use same robust path finding as above + result_folder=$(find . -maxdepth 2 -name "benchmark_${policy}" -type d | head -1) + if [ -z "$result_folder" ]; then + result_folder=$(find . -maxdepth 3 -path "*benchmark_${policy}*" -type d | head -1) + fi + + if [ -n "$result_folder" ] && [ -d "$result_folder" ]; then + json_file=$(find "$result_folder" -name "*.json" -not -name "experiment_metadata.json" | head -1) + if [ -n "$json_file" ] && [ -f "$json_file" ]; then + # Extract metrics for validation + ttft=$(jq -r '.aggregated_metrics.stats.ttft.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + e2e_latency=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + input_throughput=$(jq -r '.aggregated_metrics.stats.input_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + output_throughput=$(jq -r '.aggregated_metrics.stats.output_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + + # Check thresholds (using same values as in main workflow) + validation_status="✅" + if [ "$ttft" != "N/A" ] && [ "$ttft" != "null" ]; then + if (( $(echo "$ttft > 2.0" | bc -l 2>/dev/null || echo "0") )); then + validation_status="❌" + fi + fi + if [ "$e2e_latency" != "N/A" ] && [ "$e2e_latency" != "null" ]; then + if (( $(echo "$e2e_latency > 24.0" | bc -l 2>/dev/null || echo "0") )); then + validation_status="❌" + fi + fi + if [ "$input_throughput" != "N/A" ] && [ "$input_throughput" != "null" ]; then + if (( $(echo "$input_throughput < 10000" | bc -l 2>/dev/null || echo "0") )); then + validation_status="❌" + fi + fi + if [ "$output_throughput" != "N/A" ] && [ "$output_throughput" != "null" ]; then + if (( $(echo "$output_throughput < 90" | bc -l 2>/dev/null || echo "0") )); then + validation_status="❌" + fi + fi + + validation_summary="${validation_summary}- **${policy}**: $validation_status\n" + else + validation_summary="${validation_summary}- **${policy}**: ❌ No data\n" + fi + else + validation_summary="${validation_summary}- **${policy}**: ❌ Failed\n" + fi + done + + echo -e "$validation_summary" >> $GITHUB_STEP_SUMMARY + + echo "" >> $GITHUB_STEP_SUMMARY + echo "## 📊 Genai-Bench Features Used" >> $GITHUB_STEP_SUMMARY + echo "- **Token-level Performance**: TTFT, TPOT, End-to-End latency" >> $GITHUB_STEP_SUMMARY + echo "- **Throughput Analysis**: Input/Output/Total token throughput" >> $GITHUB_STEP_SUMMARY + echo "- **Statistical Analysis**: Percentiles, mean, std dev for all metrics" >> $GITHUB_STEP_SUMMARY + echo "- **Visual Reports**: Automated plots and Excel summaries" >> $GITHUB_STEP_SUMMARY + echo "- **SGLang Backend**: Native integration with SGLang serving" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "✅ All policies tested successfully with genai-bench!" >> $GITHUB_STEP_SUMMARY diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml new file mode 100644 index 000000000..60c7ebdf2 --- /dev/null +++ b/.github/workflows/pr-test-rust.yml @@ -0,0 +1,190 @@ +name: PR Test (Rust) + +on: + push: + branches: [ main ] + paths: + - "sgl-router/**" + pull_request: + branches: [ main ] + paths: + - "sgl-router/**" + workflow_dispatch: + +concurrency: + group: pr-test-rust-${{ github.ref }} + cancel-in-progress: true + +jobs: + unit-test-rust: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + + - name: Run lint + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + cargo clippy --all-targets --all-features -- -D warnings + + - name: Run fmt + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + cargo fmt -- --check + + - name: Run Rust tests + timeout-minutes: 20 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + cargo test + + - name: Check benchmark compilation + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + cargo check --benches + + - name: Quick benchmark sanity check + timeout-minutes: 15 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + # Run quick benchmarks to ensure they work using Python script + python3 scripts/run_benchmarks.py --quick + + pytest-rust: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + runs-on: BM.A10.4 + timeout-minutes: 25 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install rust dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Install SGLang dependencies + run: | + sudo bash scripts/ci/ci_install_dependency.sh + + - name: Build python binding + run: | + source "$HOME/.cargo/env" + cd sgl-router + pip install setuptools-rust wheel build + python3 -m build + pip install --force-reinstall dist/*.whl + + + - name: Run Python unit tests + run: | + cd sgl-router + source "$HOME/.cargo/env" + pip install pytest pytest-cov pytest-xdist + pytest -q py_test/unit --cov=sglang_router --cov-report=term-missing --cov-fail-under=80 + + - name: Run Python integration tests + run: | + cd sgl-router + source "$HOME/.cargo/env" + # Integration tests use FastAPI/uvicorn for mock workers + pip install fastapi uvicorn orjson + pytest -q -m integration + + - name: Run Python E2E tests + run: | + bash scripts/killall_sglang.sh "nuk_gpus" + cd sgl-router + python3 -m pip --no-cache-dir install --upgrade --ignore-installed blinker + python3 -m pip --no-cache-dir install --upgrade --break-system-packages genai-bench==0.0.2 + pytest -m e2e -s -vv -o log_cli=true --log-cli-level=INFO + + - name: Upload benchmark results + if: success() + uses: actions/upload-artifact@v4 + with: + name: genai-bench-results-all-policies + path: sgl-router/benchmark_**/ + + finish: + needs: [unit-test-rust, pytest-rust] + runs-on: ubuntu-latest + steps: + - name: Finish + run: echo "This is an empty step to ensure that all jobs are completed." + + summarize-benchmarks: + needs: pytest-rust + runs-on: ubuntu-latest + if: success() + + steps: + - name: Install jq + run: sudo apt-get update && sudo apt-get install -y jq bc + + - name: Download benchmark results + uses: actions/download-artifact@v4 + with: + name: genai-bench-results-all-policies + + - name: List downloaded contents + run: | + echo "Contents after download:" + ls -la + find . -name "benchmark_*" -type d + echo "JSON files found:" + find . -name "*.json" | head -10 + + - name: Create benchmark summary + run: | + echo "=== DEBUG: Creating benchmark summary ===" + echo "Available benchmark directories:" + find . -name "benchmark_*" -type d || true + echo "==========================================" + + echo "## Router E2E Genai-Bench Results Summary" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "Results captured from E2E tests for two scenarios: regular router (2 workers, dp=2) and PD router (2 prefill + 2 decode)." >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "| Scenario | Status | TTFT (s) | E2E Latency (s) | Input Throughput (tok/s) | Output Throughput (tok/s) |" >> $GITHUB_STEP_SUMMARY + echo "|----------|--------|----------|-----------------|--------------------------|---------------------------|" >> $GITHUB_STEP_SUMMARY + + scenarios=$'Regular (dp=2, round_robin)|benchmark_round_robin_regular\nPD (2 prefill + 2 decode, round_robin)|benchmark_round_robin_pd' + + echo "$scenarios" | sed 's/^\s*//' | while IFS='|' read -r label pattern; do + [ -z "$label" ] && continue + # Find the result folder (handle different extraction layouts) + result_folder=$(find . -maxdepth 3 \( -name "$pattern" -o -path "*${pattern}*" \) -type d | head -1) + + if [ -n "$result_folder" ] && [ -d "$result_folder" ]; then + json_file=$(find "$result_folder" -name "*.json" -not -name "experiment_metadata.json" | head -1) + + if [ -n "$json_file" ] && [ -f "$json_file" ]; then + ttft_mean=$(jq -r '.aggregated_metrics.stats.ttft.mean' "$json_file") + e2e_latency_mean=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean' "$json_file") + input_throughput_mean=$(jq -r '.aggregated_metrics.stats.input_throughput.mean' "$json_file") + output_throughput_mean=$(jq -r '.aggregated_metrics.stats.output_throughput.mean' "$json_file") + + ttft_display=$(printf "%.2f" "$ttft_mean" 2>/dev/null || echo "$ttft_mean") + e2e_display=$(printf "%.2f" "$e2e_latency_mean" 2>/dev/null || echo "$e2e_latency_mean") + input_display=$(printf "%.0f" "$input_throughput_mean" 2>/dev/null || echo "$input_throughput_mean") + output_display=$(printf "%.0f" "$output_throughput_mean" 2>/dev/null || echo "$output_throughput_mean") + + echo "| ${label} | ✅ Success | $ttft_display | $e2e_display | $input_display | $output_display |" >> $GITHUB_STEP_SUMMARY + fi + fi + done diff --git a/.github/workflows/pr-test-sgl-kernel.yml b/.github/workflows/pr-test-sgl-kernel.yml new file mode 100644 index 000000000..832188cdd --- /dev/null +++ b/.github/workflows/pr-test-sgl-kernel.yml @@ -0,0 +1,151 @@ +name: PR Test (sgl-kernel) + +on: + push: + branches: [main] + paths: + - "sgl-kernel/**" + pull_request: + branches: [main] + paths: + - "sgl-kernel/**" + workflow_dispatch: + +concurrency: + group: pr-test-sgl-kernel-${{ github.ref }} + cancel-in-progress: true + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Check clang-format + uses: DoozyX/clang-format-lint-action@v0.18.1 + with: + source: sgl-kernel + extensions: h,c,cpp,hpp,cu,cuh,cc + clangFormatVersion: 18 + style: file + + build-wheels: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + runs-on: sgl-kernel-build-node + strategy: + matrix: + include: + - python-version: "3.10" + cuda-version: "12.4" + - python-version: "3.10" + cuda-version: "12.8" + - python-version: "3.10" + cuda-version: "12.9" + name: Build Wheel (CUDA ${{ matrix.cuda-version }}) + steps: + - name: Cleanup + run: | + sudo rm -rf $GITHUB_WORKSPACE/* || true + + - uses: actions/checkout@v4 + with: + submodules: "recursive" + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Build wheel for Python ${{ matrix.python-version }} and CUDA ${{ matrix.cuda-version }} + if: github.event_name != 'push' || (matrix.cuda-version != '12.4' && matrix.cuda-version != '12.8') + run: | + cd sgl-kernel + chmod +x ./build.sh + ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" + + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }} + path: sgl-kernel/dist/* + + unit-test: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + needs: build-wheels + runs-on: 1-gpu-runner + steps: + - uses: actions/checkout@v4 + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + path: sgl-kernel/dist/ + merge-multiple: true + pattern: wheel-python3.10-cuda12.9 + + - name: Install + run: | + bash scripts/ci/ci_install_dependency.sh + pip3 install torch==2.8.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu126 && pip3 install pytest + pip3 uninstall sgl-kernel -y || true + pip3 install sgl-kernel/dist/*whl --force-reinstall --no-deps + pip3 list | grep sgl-kernel + + - name: Run test + timeout-minutes: 30 + run: | + cd sgl-kernel + pytest tests/ + + - name: Uninstall dependencies + run: | + pip3 uninstall sgl-kernel -y + + mla-test: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + needs: build-wheels + runs-on: 1-gpu-runner + steps: + - uses: actions/checkout@v4 + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + path: sgl-kernel/dist/ + merge-multiple: true + pattern: wheel-python3.10-cuda12.9 + + - name: Install + run: | + bash scripts/ci/ci_install_dependency.sh + pip3 install torch==2.8.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu126 + pip3 uninstall sgl-kernel -y || true + pip3 install sgl-kernel/dist/*whl --force-reinstall --no-deps + pip3 list | grep sgl-kernel + + - name: Run test + timeout-minutes: 30 + run: | + cd test/srt + python3 test_mla_deepseek_v3.py + + - name: Uninstall dependencies + run: | + pip3 uninstall sgl-kernel -y + + finish: + needs: [unit-test, mla-test, lint, build-wheels] + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + results=(${{ join(needs.*.result, ' ') }}) + for result in "${results[@]}"; do + if [ "$result" = "failure" ] || [ "$result" = "cancelled" ]; then + echo "Job failed with result: $result" + exit 1 + fi + done + echo "All jobs completed successfully" + exit 0 diff --git a/.github/workflows/pr-test-xeon.yml b/.github/workflows/pr-test-xeon.yml new file mode 100644 index 000000000..fcc70f286 --- /dev/null +++ b/.github/workflows/pr-test-xeon.yml @@ -0,0 +1,106 @@ +name: PR Test (Xeon) + +on: + push: + branches: [ main ] + paths: + - "python/**" + - "scripts/ci/**" + - "test/**" + - "sgl-kernel/**" + - ".github/workflows/pr-test-xeon.yml" + pull_request: + branches: [ main ] + paths: + - "python/**" + - "scripts/ci/**" + - "test/**" + - "sgl-kernel/**" + - ".github/workflows/pr-test-xeon.yml" + workflow_dispatch: + +concurrency: + group: pr-test-xeon-${{ github.ref }} + cancel-in-progress: false + +jobs: + build-test: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + runs-on: xeon-gnr + env: + HF_HOME: /home/sdp/.cache/huggingface + strategy: + matrix: + build_type: ['all'] + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Build and Push + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + tag=v${version}-xeon + + docker build . -f docker/Dockerfile.xeon -t sglang_xeon --no-cache + + - name: Run container + run: | + docker run -dt \ + -v ${{ github.workspace }}:/sglang-checkout/ --ipc=host \ + -v ${HF_HOME}:/root/.cache/huggingface \ + --name ci_sglang_xeon \ + sglang_xeon + + - name: Install dependencies + timeout-minutes: 20 + run: | + docker exec ci_sglang_xeon bash -c "python3 -m pip install --upgrade pip" + docker exec ci_sglang_xeon pip uninstall sgl-kernel -y || true + docker exec -w /sglang-checkout/sgl-kernel ci_sglang_xeon bash -c "cp pyproject_cpu.toml pyproject.toml && pip install -v ." + docker exec -w /sglang-checkout/ ci_sglang_xeon bash -c "pip install -e "python[dev_cpu]"" + + - name: Check AMX support + id: check_amx + timeout-minutes: 5 + run: | + docker exec -w /sglang-checkout/ ci_sglang_xeon \ + bash -c "python3 -c 'import torch; import sgl_kernel; assert torch._C._cpu._is_amx_tile_supported(); assert hasattr(torch.ops.sgl_kernel, \"convert_weight_packed\"); '" + continue-on-error: true + + - name: Run unit tests + if: steps.check_amx.outcome == 'success' + timeout-minutes: 36 + run: | + docker exec -w /sglang-checkout/ ci_sglang_xeon \ + bash -c "cd ./test/srt && python3 run_suite.py --suite per-commit-cpu" + + - name: Change permission + timeout-minutes: 2 + run: | + docker exec -u root ci_sglang_xeon bash -c " + rm -rf /tmp/ci-home && + chown -R $(id -u):$(id -g) /sglang-checkout/ 2>/dev/null || true + " + + - name: Cleanup container + if: always() + run: | + docker rm -f ci_sglang_xeon || true + + pr-test-xeon-finish: + if: always() + needs: [build-test] + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + results=(${{ join(needs.*.result, ' ') }}) + for result in "${results[@]}"; do + if [ "$result" = "failure" ] || [ "$result" = "cancelled" ]; then + echo "Job failed with result: $result" + exit 1 + fi + done + echo "All jobs completed successfully" + exit 0 diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml new file mode 100644 index 000000000..bd1053902 --- /dev/null +++ b/.github/workflows/pr-test.yml @@ -0,0 +1,437 @@ +name: PR Test + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + inputs: + version: + description: "FlashInfer version" + required: true + type: choice + default: 'release' + options: + - 'release' + - 'nightly' + +concurrency: + group: pr-test-${{ github.ref }} + cancel-in-progress: true + +jobs: + check-changes: + runs-on: ubuntu-latest + outputs: + src: ${{ steps.filter.outputs.src }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + with: + filters: | + src: + - "python/**" + - "scripts/ci/**" + - "test/**" + - ".github/workflows/pr-test.yml" + + unit-test-frontend: + needs: check-changes + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 10 + run: | + cd test/lang + python3 run_suite.py --suite per-commit + + unit-test-backend-1-gpu: + needs: [check-changes, unit-test-frontend] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 1-gpu-runner + strategy: + fail-fast: false + matrix: + part: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 30 + run: | + cd test/srt + python3 run_suite.py --suite per-commit --auto-partition-id ${{ matrix.part }} --auto-partition-size 10 + + unit-test-backend-2-gpu: + needs: [check-changes] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 2-gpu-runner + strategy: + fail-fast: false + matrix: + part: [0, 1] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 30 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-2-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 + + unit-test-backend-4-gpu: + needs: [check-changes, unit-test-backend-2-gpu] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 4-gpu-runner + strategy: + fail-fast: false + matrix: + part: [0, 1] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-4-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 + + unit-test-backend-8-gpu: + needs: [check-changes, unit-test-backend-2-gpu] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 8-gpu-runner + strategy: + fail-fast: false + matrix: + part: [0, 1] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-8-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 + + performance-test-1-gpu-part-1: + needs: check-changes + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Benchmark single latency + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_bs1_small + python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_bs1_default + + - name: Benchmark online latency + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_online_latency_default + + - name: Benchmark offline throughput + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_default + + - name: Benchmark offline throughput (Non-streaming, small batch size) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_non_stream_small_batch_size + + - name: Benchmark online latency (EAGLE) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_online_latency_eagle + + - name: Benchmark online latency (LoRA) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_lora_online_latency + python3 -m unittest test_bench_serving.TestBenchServing.test_lora_online_latency_with_concurrent_adapter_updates + + performance-test-1-gpu-part-2: + needs: check-changes + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Benchmark offline throughput (w/o RadixAttention) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_without_radix_cache + + - name: Benchmark offline throughput (w/ Triton) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_with_triton_attention_backend + + - name: Benchmark offline throughput (w/ FP8) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_default_fp8 + + - name: Benchmark VLM offline throughput + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_vlm_offline_throughput + + - name: Benchmark VLM online latency + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_vlm_online_latency + + performance-test-2-gpu: + needs: [check-changes, unit-test-backend-2-gpu] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 2-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + + - name: Benchmark single latency (TP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_moe_tp2_bs1 + + - name: Benchmark single latency + torch.compile (TP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_torch_compile_tp2_bs1 + + - name: Benchmark offline throughput (TP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_default + + - name: Benchmark offline throughput (w/o RadixAttention) (TP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_without_radix_cache + + - name: Benchmark offline PP decode throughput (PP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_pp_offline_throughput_default_decode + + - name: Benchmark offline PP prefill throughput (PP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_pp_long_context_prefill + + accuracy-test-1-gpu: + needs: check-changes + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + git clone https://github.com/merrymercy/human-eval.git + cd human-eval + pip install -e . + + - name: Evaluate accuracy + timeout-minutes: 20 + run: | + cd test/srt + python3 test_eval_accuracy_large.py + + accuracy-test-2-gpu: + needs: [check-changes, accuracy-test-1-gpu] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 2-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + git clone https://github.com/merrymercy/human-eval.git + cd human-eval + pip install -e . + + - name: Evaluate accuracy (TP=2) + timeout-minutes: 20 + run: | + cd test/srt + python3 test_moe_eval_accuracy_large.py + + unit-test-deepep-4-gpu: + needs: [check-changes, unit-test-backend-2-gpu] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 4-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_deepep.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-4-gpu-deepep + + unit-test-deepep-8-gpu: + needs: [check-changes, unit-test-backend-2-gpu] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: 8-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_deepep.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-8-gpu-deepep + + unit-test-backend-8-gpu-b200: + needs: [check-changes, unit-test-backend-2-gpu] + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false && + needs.check-changes.outputs.src == 'true' + runs-on: b200-runner + strategy: + fail-fast: false + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + IS_BLACKWELL=1 bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-8-gpu-b200 --auto-partition-id 0 --auto-partition-size 1 + + + pr-test-finish: + needs: [ + check-changes, + unit-test-frontend, unit-test-backend-1-gpu, + unit-test-backend-2-gpu, unit-test-backend-4-gpu, unit-test-backend-8-gpu, + performance-test-1-gpu-part-1, performance-test-1-gpu-part-2, performance-test-2-gpu, + accuracy-test-1-gpu, accuracy-test-2-gpu, + unit-test-deepep-4-gpu, unit-test-deepep-8-gpu, + unit-test-backend-8-gpu-b200, + ] + if: always() + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + results=(${{ join(needs.*.result, ' ') }}) + for result in "${results[@]}"; do + if [ "$result" = "failure" ] || [ "$result" = "cancelled" ]; then + echo "Job failed with result: $result" + exit 1 + fi + done + echo "All jobs completed successfully" + exit 0 diff --git a/.github/workflows/release-docker-amd-nightly.yml b/.github/workflows/release-docker-amd-nightly.yml new file mode 100644 index 000000000..c61e200df --- /dev/null +++ b/.github/workflows/release-docker-amd-nightly.yml @@ -0,0 +1,65 @@ +name: Release Docker Images Nightly (AMD) +on: + workflow_dispatch: + schedule: + - cron: '0 13 * * *' + +concurrency: + # A PR number if a pull request and otherwise the commit hash. This cancels + # queued and in-progress runs for the same PR (presubmit) or commit + # (postsubmit). The workflow name is prepended to avoid conflicts between + # different workflows. + group: ${{ github.workflow }}-${{ github.event.number || github.sha }} + cancel-in-progress: true + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: amd-docker-scale + environment: 'prod' + strategy: + matrix: + gpu_arch: ['gfx942', 'gfx942-rocm700', 'gfx950'] + build_type: ['all', 'srt'] + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: "Set Date" + run: | + echo "DATE=$(date +%Y%m%d)" >> $GITHUB_ENV + + - name: Login to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_AMD_USERNAME }} + password: ${{ secrets.DOCKERHUB_AMD_TOKEN }} + + - name: Build and Push + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + + if [ "${{ matrix.gpu_arch }}" = "gfx942" ]; then + rocm_tag="rocm630-mi30x" + elif [ "${{ matrix.gpu_arch }}" = "gfx942-rocm700" ]; then + rocm_tag="rocm700-mi30x" + elif [ "${{ matrix.gpu_arch }}" = "gfx950" ]; then + rocm_tag="rocm700-mi35x" + else + echo "Unsupported gfx arch" + exit 1 + fi + + tag=v${version}-${rocm_tag} + + if [ "${{ matrix.build_type }}" = "all" ]; then + tag_suffix="" + elif [ "${{ matrix.build_type }}" = "srt" ]; then + tag_suffix="-srt" + else + echo "Unsupported build type" + exit 1 + fi + + docker build . -f docker/Dockerfile.rocm --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} -t rocm/sgl-dev:${tag}-${{ env.DATE }}${tag_suffix} --no-cache + docker push rocm/sgl-dev:${tag}-${{ env.DATE }}${tag_suffix} diff --git a/.github/workflows/release-docker-amd.yml b/.github/workflows/release-docker-amd.yml new file mode 100644 index 000000000..98c11e2fa --- /dev/null +++ b/.github/workflows/release-docker-amd.yml @@ -0,0 +1,56 @@ +name: Release Docker Images (AMD) +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: amd-docker-scale + environment: 'prod' + strategy: + matrix: + gpu_arch: ['gfx942', 'gfx942-rocm700', 'gfx950'] + build_type: ['all', 'srt'] + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Login to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Build and Push + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + + if [ "${{ matrix.gpu_arch }}" = "gfx942" ]; then + rocm_tag="rocm630-mi30x" + elif [ "${{ matrix.gpu_arch }}" = "gfx942-rocm700" ]; then + rocm_tag="rocm700-mi30x" + elif [ "${{ matrix.gpu_arch }}" = "gfx950" ]; then + rocm_tag="rocm700-mi35x" + else + echo "Unsupported gfx arch" + exit 1 + fi + + tag=v${version}-${rocm_tag} + + if [ "${{ matrix.build_type }}" = "all" ]; then + tag_suffix="" + elif [ "${{ matrix.build_type }}" = "srt" ]; then + tag_suffix="-srt" + else + echo "Unsupported build type" + exit 1 + fi + + docker build . -f docker/Dockerfile.rocm --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} -t lmsysorg/sglang:${tag}${tag_suffix} --no-cache + docker push lmsysorg/sglang:${tag}${tag_suffix} diff --git a/.github/workflows/release-docker-dev.yml b/.github/workflows/release-docker-dev.yml new file mode 100644 index 000000000..38e2e790f --- /dev/null +++ b/.github/workflows/release-docker-dev.yml @@ -0,0 +1,49 @@ +name: Build Development Docker Image + +on: + workflow_dispatch: + schedule: + - cron: '0 0 * * *' + +jobs: + build-dev: + if: ${{ github.repository == 'sgl-project/sglang' }} + runs-on: ubuntu-22.04 + strategy: + matrix: + variant: + - version: 12.6.1 + type: all + tag: dev + - version: 12.8.1 + type: blackwell + tag: blackwell + - version: 12.9.1 + type: blackwell + tag: b200-cu129 + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Free disk space + uses: jlumbroso/free-disk-space@main + with: + tool-cache: false + docker-images: false + android: true + dotnet: true + haskell: true + large-packages: true + swap-storage: false + + - name: Login to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Build and Push Dev Image + run: | + docker buildx build --output type=image,compression=zstd . -f docker/Dockerfile --build-arg CUDA_VERSION=${{ matrix.variant.version }} --build-arg BUILD_TYPE=${{ matrix.variant.type }} --build-arg CMAKE_BUILD_PARALLEL_LEVEL=$(nproc) -t lmsysorg/sglang:${{ matrix.variant.tag }} --no-cache + docker push lmsysorg/sglang:${{ matrix.variant.tag }} diff --git a/.github/workflows/release-docker-gb200.yml b/.github/workflows/release-docker-gb200.yml new file mode 100644 index 000000000..fbcacb330 --- /dev/null +++ b/.github/workflows/release-docker-gb200.yml @@ -0,0 +1,36 @@ +name: Release Docker Images (GB200) +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-22.04-arm + environment: "prod" + steps: + - name: Delete huge unnecessary tools folder + run: rm -rf /opt/hostedtoolcache + + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Build and Push + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + tag=v${version}-cu129-gb200 + + docker buildx build --platform linux/arm64 --push --output type=image -t lmsysorg/sglang:${tag} -f docker/Dockerfile.gb200 --build-arg CUDA_VERSION=12.9.1 --build-arg BUILD_TYPE=blackwell --no-cache . diff --git a/.github/workflows/release-docker-npu-nightly.yml b/.github/workflows/release-docker-npu-nightly.yml new file mode 100644 index 000000000..527a0cdc2 --- /dev/null +++ b/.github/workflows/release-docker-npu-nightly.yml @@ -0,0 +1,78 @@ +name: Release Docker Images Nightly (Ascend NPU) +on: + pull_request: + branches: + - main + paths: + - ".github/workflows/release-docker-npu-nightly.yml" + - "docker/Dockerfile.npu" + workflow_dispatch: + schedule: + - cron: "0 0 * * *" + +concurrency: + group: ${{ github.workflow }}-${{ github.sha }} + cancel-in-progress: true + +jobs: + build: + runs-on: ubuntu-22.04-arm + strategy: + matrix: + cann_version: ["8.2.rc1"] + device_type: ["910b", "a3"] + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Free up disk space + uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be # v1.3.1 + with: + tool-cache: true + docker-images: false + + - name: Setup Docker buildx + uses: docker/setup-buildx-action@v3 + + - name: Docker meta + id: meta + uses: docker/metadata-action@v5 + with: + images: | + lmsysorg/sglang + # push with schedule event + # push with workflow_dispatch event + tags: | + type=ref,event=pr + type=ref,event=branch + type=schedule,pattern=main + flavor: | + latest=false + suffix=-cann${{ matrix.cann_version }}-${{ matrix.device_type }},onlatest=true + # Login against a Docker registry except on PR + # https://github.com/docker/login-action + - name: Log into docker hub + uses: docker/login-action@v3 + if: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + # Build and push Docker image with Buildx (don't push on PR) + # https://github.com/docker/build-push-action + - name: Build and push Docker image + id: build-and-push + uses: docker/build-push-action@v6 + with: + context: docker + file: docker/Dockerfile.npu + # TODO: need add x86 platforms support when memfabric is ready + platforms: linux/arm64 + labels: ${{ steps.meta.outputs.labels }} + tags: ${{ steps.meta.outputs.tags }} + push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} + provenance: false + build-args: | + SGLANG_KERNEL_NPU_TAG=20250901 + CANN_VERSION=${{ matrix.cann_version }} + DEVICE_TYPE=${{ matrix.device_type }} diff --git a/.github/workflows/release-docker-npu.yml b/.github/workflows/release-docker-npu.yml new file mode 100644 index 000000000..f9d52eb4b --- /dev/null +++ b/.github/workflows/release-docker-npu.yml @@ -0,0 +1,74 @@ +name: Release Docker Images (Ascend NPU) +on: + push: + tags: + - "*" # Trigger on all tags and filterred by pep440 later + workflow_dispatch: + pull_request: + branches: + - main + paths: + - ".github/workflows/release-docker-npu.yml" + - "docker/Dockerfile.npu" + +jobs: + build: + runs-on: ubuntu-22.04-arm + strategy: + matrix: + cann_version: ["8.2.rc1"] + device_type: ["910b", "a3"] + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Free up disk space + uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be # v1.3.1 + with: + tool-cache: true + docker-images: false + + # push with tag + - name: Docker meta + id: meta + uses: docker/metadata-action@v5 + with: + images: | + lmsysorg/sglang + tags: | + type=ref,event=pr + type=ref,event=tag,suffix=-cann${{ matrix.cann_version }}-${{ matrix.device_type }} + flavor: | + latest=false + + # Login against a Docker registry except on PR + # https://github.com/docker/login-action + - name: Login to Docker Hub + uses: docker/login-action@v2 + if: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Get version + id: get_version + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + echo "TAG=lmsysorg/sglang:v$version-cann${{ matrix.cann_version }}-${{ matrix.device_type }}" >> $GITHUB_OUTPUT + + - name: Build and push Docker image + id: build-and-push + uses: docker/build-push-action@v6 + with: + context: docker + file: docker/Dockerfile.npu + # TODO: need add x86 platforms support when memfabric is ready + platforms: linux/arm64 + labels: ${{ steps.meta.outputs.labels }} + tags: ${{ steps.meta.outputs.tags || steps.get_version.outputs.TAG }} + push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} + provenance: false + build-args: | + SGLANG_KERNEL_NPU_TAG=20250901 + CANN_VERSION=${{ matrix.cann_version }} + DEVICE_TYPE=${{ matrix.device_type }} diff --git a/.github/workflows/release-docker-router.yml b/.github/workflows/release-docker-router.yml new file mode 100644 index 000000000..f98651e8a --- /dev/null +++ b/.github/workflows/release-docker-router.yml @@ -0,0 +1,30 @@ +name: Release SGLang Router Docker Image +on: + push: + branches: + - main + paths: + - "sgl-router/py_src/sglang_router/version.py" + workflow_dispatch: + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-24.04 + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Login to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Build and Push + run: | + version=$(cat sgl-router/py_src/sglang_router/version.py | cut -d'"' -f2) + tag=v${version} + + docker build . -f docker/Dockerfile.router -t lmsysorg/sglang-router:${tag} --no-cache + docker push lmsysorg/sglang-router:${tag} diff --git a/.github/workflows/release-docker-xeon.yml b/.github/workflows/release-docker-xeon.yml new file mode 100644 index 000000000..bd2a3910f --- /dev/null +++ b/.github/workflows/release-docker-xeon.yml @@ -0,0 +1,35 @@ +name: Release Docker Xeon Images +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-24.04 + environment: 'prod' + strategy: + matrix: + build_type: ['all'] + steps: + + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Login to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Build and Push + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + tag=v${version}-xeon + + docker build . -f docker/Dockerfile.xeon -t lmsysorg/sglang:${tag} --no-cache + docker push lmsysorg/sglang:${tag} diff --git a/.github/workflows/release-docker.yml b/.github/workflows/release-docker.yml new file mode 100644 index 000000000..60a8df621 --- /dev/null +++ b/.github/workflows/release-docker.yml @@ -0,0 +1,97 @@ +name: Release Docker Images +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-latest + environment: 'prod' + strategy: + matrix: + cuda_version: ['12.6.1', '12.8.1', '12.9.1'] + build_type: ['all', 'blackwell'] + exclude: + - cuda_version: '12.6.1' + build_type: 'blackwell' + - cuda_version: '12.8.1' + build_type: 'all' + - cuda_version: '12.9.1' + build_type: 'all' + steps: + - name: Delete huge unnecessary tools folder + run: rm -rf /opt/hostedtoolcache + + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Free disk space + uses: jlumbroso/free-disk-space@main + with: + tool-cache: false + docker-images: false + android: true + dotnet: true + haskell: true + large-packages: true + swap-storage: false + + - name: Login to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Build and Push + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + + if [ "${{ matrix.cuda_version }}" = "11.8.0" ]; then + cuda_tag="cu118" + elif [ "${{ matrix.cuda_version }}" = "12.1.1" ]; then + cuda_tag="cu121" + elif [ "${{ matrix.cuda_version }}" = "12.4.1" ]; then + cuda_tag="cu124" + elif [ "${{ matrix.cuda_version }}" = "12.5.1" ]; then + cuda_tag="cu125" + elif [ "${{ matrix.cuda_version }}" = "12.6.1" ]; then + cuda_tag="cu126" + elif [ "${{ matrix.cuda_version }}" = "12.8.1" ]; then + cuda_tag="cu128" + elif [ "${{ matrix.cuda_version }}" = "12.9.1" ]; then + cuda_tag="cu129" + else + echo "Unsupported CUDA version" + exit 1 + fi + + tag=v${version}-${cuda_tag} + + if [ "${{ matrix.build_type }}" = "all" ]; then + tag_suffix="" + elif [ "${{ matrix.build_type }}" = "srt" ]; then + tag_suffix="-srt" + elif [ "${{ matrix.build_type }}" = "blackwell" ]; then + tag_suffix="-b200" + else + echo "Unsupported build type" + exit 1 + fi + + docker buildx build --output type=image,compression=zstd . -f docker/Dockerfile --build-arg CUDA_VERSION=${{ matrix.cuda_version }} --build-arg BUILD_TYPE=${{ matrix.build_type }} -t lmsysorg/sglang:${tag}${tag_suffix} --no-cache + docker push lmsysorg/sglang:${tag}${tag_suffix} + + if [ "${{ matrix.cuda_version }}" = "12.6.1" ]; then + docker tag lmsysorg/sglang:${tag}${tag_suffix} lmsysorg/sglang:latest${tag_suffix} + docker push lmsysorg/sglang:latest${tag_suffix} + fi + + if [ "${{ matrix.cuda_version }}" = "12.9.1" ]; then + docker tag lmsysorg/sglang:${tag}${tag_suffix} lmsysorg/sglang:v${version} + docker push lmsysorg/sglang:v${version} + fi diff --git a/.github/workflows/release-docs.yml b/.github/workflows/release-docs.yml new file mode 100644 index 000000000..0e09eec93 --- /dev/null +++ b/.github/workflows/release-docs.yml @@ -0,0 +1,65 @@ +name: Release Documentation + +on: + push: + branches: + - main + paths: + - "docs/**" + - "python/sglang/version.py" + - "python/sglang/**" + workflow_dispatch: + +concurrency: + group: release-docs-${{ github.ref }} + cancel-in-progress: true + +jobs: + execute-and-deploy: + runs-on: 1-gpu-runner + if: github.repository == 'sgl-project/sglang' + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + pip install -r docs/requirements.txt + apt-get update && apt-get install -y pandoc parallel retry + ln -sf "$(which python3)" /usr/bin/python + + - name: Setup Jupyter Kernel + run: | + python -m ipykernel install --user --name python3 --display-name "Python 3" + + - name: Execute notebooks + timeout-minutes: 40 + run: | + cd docs + make clean + make compile + + - name: Push HTML to sgl-project.github.io + timeout-minutes: 60 + env: + GITHUB_TOKEN: ${{ secrets.DOCUMENTATION_PAT_TOKEN }} + run: | + cd docs + make html + python3 wrap_run_llm.py + + cd _build/html + + git clone https://$GITHUB_TOKEN@github.com/sgl-project/sgl-project.github.io.git ../sgl-project.github.io --depth 1 + find ../sgl-project.github.io/ -mindepth 1 -not -path "../sgl-project.github.io/.git*" -not -name CNAME -not -name ".jekyll" -not -name ".nojekyll" -delete + cp -r * ../sgl-project.github.io + cp ../../README.md ../sgl-project.github.io/README.md + cd ../sgl-project.github.io + git config user.name "zhaochenyang20" + git config user.email "zhaochenyang20@gmail.com" + git add . + git commit -m "Update $(date +'%Y-%m-%d %H:%M:%S')" + git push https://$GITHUB_TOKEN@github.com/sgl-project/sgl-project.github.io.git main + cd .. + rm -rf sgl-project.github.io diff --git a/.github/workflows/release-fake-tag.yml b/.github/workflows/release-fake-tag.yml new file mode 100644 index 000000000..ce5999506 --- /dev/null +++ b/.github/workflows/release-fake-tag.yml @@ -0,0 +1,35 @@ +name: Release Fake Tag +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +permissions: + contents: write + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-latest + environment: 'prod' + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Get version + id: get_version + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + echo "TAG=v$version" >> $GITHUB_OUTPUT + + - name: Create and push fake tag + env: + GITHUB_TOKEN: ${{ secrets.REPO_TOKEN }} + run: | + git config user.name zhyncs + git config user.email me@zhyncs.com + git checkout -b ${{ steps.get_version.outputs.TAG }} + git push --set-upstream origin ${{ steps.get_version.outputs.TAG }} diff --git a/.github/workflows/release-pypi-router.yml b/.github/workflows/release-pypi-router.yml new file mode 100644 index 000000000..a2128be83 --- /dev/null +++ b/.github/workflows/release-pypi-router.yml @@ -0,0 +1,119 @@ +# Reference: https://github.com/openai/tiktoken/blob/63527649963def8c759b0f91f2eb69a40934e468/.github/workflows/build_wheels.yml#L1 + +name: Release SGLang Router to PyPI + +on: + push: + branches: + - main + paths: + - sgl-router/pyproject.toml + workflow_dispatch: + +jobs: + build: + name: Build on ${{ matrix.os }} (${{ matrix.target }}) + runs-on: ${{ matrix.os }}-latest + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu + target: x86_64 + + steps: + - uses: actions/checkout@v4 + with: + path: sglang-repo + + - name: Move sgl-router folder to root and delete sglang-repo + run: | + mv sglang-repo/sgl-router/* . + rm -rf sglang-repo + ls -alt + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install build dependencies + run: | + python -m pip install -U pip + python -m pip install build twine auditwheel + + - name: Build package + uses: pypa/cibuildwheel@v2.21.3 + env: + CIBW_BUILD: "cp38-manylinux_x86_64 cp39-manylinux_x86_64 cp310-manylinux_x86_64 cp311-manylinux_x86_64 cp312-manylinux_x86_64" + CIBW_BEFORE_ALL: | + yum update -y && yum install -y openssl-devel wget unzip && \ + # Install latest protoc (v32.0) that supports proto3 + cd /tmp && \ + wget https://github.com/protocolbuffers/protobuf/releases/download/v32.0/protoc-32.0-linux-x86_64.zip && \ + unzip protoc-32.0-linux-x86_64.zip -d /usr/local && \ + rm protoc-32.0-linux-x86_64.zip && \ + # Install Rust + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + CIBW_ENVIRONMENT: "PATH=$HOME/.cargo/bin:$PATH" + + - name: List built packages + run: ls -lh wheelhouse/ + + - name: Check packages + run: twine check --strict wheelhouse/* + + - uses: actions/upload-artifact@v4 + with: + name: packages-${{ matrix.os }}-${{ matrix.target }} + path: wheelhouse/ + + build-sdist: + name: Build SDist + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + path: sglang-repo + + - name: Move sgl-router folder to root, copy the license file, and delete sglang-repo + run: | + mv sglang-repo/sgl-router/* . + mv sglang-repo/LICENSE . + rm -rf sglang-repo + ls -alt + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Build SDist + run: | + pip install build + python -m pip install -U packaging + python -m build --sdist + + - uses: actions/upload-artifact@v4 + with: + name: sdist + path: dist/*.tar.gz + + upload: + name: Upload to PyPI + if: github.repository == 'sgl-project/sglang' # Ensure this job only runs for the sgl-project/sglang repository + needs: [build, build-sdist] + runs-on: ubuntu-latest + steps: + - uses: actions/download-artifact@v4 + with: + path: dist + merge-multiple: true + + - name: Upload to PyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN_ROUTER }} + run: | + pip install twine + twine upload dist/* --verbose diff --git a/.github/workflows/release-pypi.yml b/.github/workflows/release-pypi.yml new file mode 100644 index 000000000..354ee21d1 --- /dev/null +++ b/.github/workflows/release-pypi.yml @@ -0,0 +1,31 @@ +name: Release PyPI +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-latest + environment: "prod" + steps: + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Upload to pypi + run: | + cd python + cp ../README.md ../LICENSE . + pip install build + python3 -m build + pip install twine + python3 -m twine upload dist/* -u __token__ -p ${{ secrets.PYPI_TOKEN }} diff --git a/.github/workflows/release-whl-kernel-cu118.yml b/.github/workflows/release-whl-kernel-cu118.yml new file mode 100644 index 000000000..4757bcaa1 --- /dev/null +++ b/.github/workflows/release-whl-kernel-cu118.yml @@ -0,0 +1,92 @@ +name: Release SGLang Kernel Wheel (cu118) + +on: + workflow_dispatch: + inputs: + tag_name: + type: string + push: + branches: + - main + paths: + - sgl-kernel/python/sgl_kernel/version.py + +jobs: + build-wheels: + if: github.repository == 'sgl-project/sglang' + runs-on: sgl-kernel-release-node + strategy: + matrix: + python-version: ["3.9"] + cuda-version: ["11.8"] + + steps: + - uses: actions/checkout@v4 + with: + submodules: "recursive" + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Build wheels for Python ${{ matrix.python-version }} and CUDA ${{ matrix.cuda-version }} + run: | + cd sgl-kernel + chmod +x ./build.sh + ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" + + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }} + path: sgl-kernel/dist/* + + release: + needs: build-wheels + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + path: sgl-kernel/dist/ + merge-multiple: true + pattern: wheel-* + + - name: Set tag name + id: set_tag_name + run: | + if [ -z "${{ inputs.tag_name }}" ]; then + TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)" + echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT + else + echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT + fi + + - name: Release + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ steps.set_tag_name.outputs.tag_name }} + repository: sgl-project/whl + token: ${{ secrets.WHL_TOKEN }} + files: | + sgl-kernel/dist/* + + - name: Clone wheel index + run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl + env: + WHL_TOKEN: ${{ secrets.WHL_TOKEN }} + + - name: Update wheel index + run: python3 scripts/update_kernel_whl_index.py + + - name: Push wheel index + run: | + cd sgl-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 index" + git push diff --git a/.github/workflows/release-whl-kernel.yml b/.github/workflows/release-whl-kernel.yml new file mode 100644 index 000000000..b12c91288 --- /dev/null +++ b/.github/workflows/release-whl-kernel.yml @@ -0,0 +1,283 @@ +name: Release SGLang Kernels + +on: + push: + branches: + - main + paths: + - sgl-kernel/python/sgl_kernel/version.py + workflow_dispatch: + inputs: + tag_name: + type: string + required: false + +concurrency: + group: release-sglang-kernels-${{ github.ref }} + cancel-in-progress: true + +jobs: + build-cu129: + if: github.repository == 'sgl-project/sglang' + runs-on: sgl-kernel-release-node + strategy: + matrix: + python-version: ["3.10"] + cuda-version: ["12.9"] + steps: + - uses: actions/checkout@v4 + with: + submodules: "recursive" + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Build wheels + run: | + cd sgl-kernel + chmod +x ./build.sh + ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" + + - name: Upload to PyPI + working-directory: sgl-kernel + run: | + pip install twine + python3 -m twine upload dist/* -u __token__ -p ${{ secrets.PYPI_TOKEN }} + + build-cu124: + if: github.repository == 'sgl-project/sglang' + needs: build-cu129 + runs-on: sgl-kernel-release-node + strategy: + matrix: + python-version: ["3.10"] + cuda-version: ["12.4"] + steps: + - uses: actions/checkout@v4 + with: + submodules: "recursive" + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Build wheels + run: | + cd sgl-kernel + chmod +x ./build.sh + ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" + + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }} + path: sgl-kernel/dist/* + + release-cu124: + needs: build-cu124 + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + path: sgl-kernel/dist/ + merge-multiple: true + pattern: wheel-* + + - name: Set tag name + id: set_tag_name + run: | + if [ -z "${{ inputs.tag_name }}" ]; then + TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)" + echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT + else + echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT + fi + + - name: Release + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ steps.set_tag_name.outputs.tag_name }} + repository: sgl-project/whl + token: ${{ secrets.WHL_TOKEN }} + files: | + sgl-kernel/dist/* + + - name: Clone wheel index + run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl + env: + WHL_TOKEN: ${{ secrets.WHL_TOKEN }} + + - name: Update wheel index + run: python3 scripts/update_kernel_whl_index.py --cuda 124 + + - name: Push wheel index + run: | + cd sgl-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 index" + git push + + build-cu128: + if: github.repository == 'sgl-project/sglang' + needs: build-cu129 + runs-on: sgl-kernel-release-node + strategy: + matrix: + python-version: ["3.10"] + cuda-version: ["12.8"] + steps: + - uses: actions/checkout@v4 + with: + submodules: "recursive" + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Build wheels + run: | + cd sgl-kernel + chmod +x ./build.sh + ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" + + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }} + path: sgl-kernel/dist/* + + release-cu128: + needs: build-cu128 + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + path: sgl-kernel/dist/ + merge-multiple: true + pattern: wheel-* + + - name: Set tag name + id: set_tag_name + run: | + if [ -z "${{ inputs.tag_name }}" ]; then + TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)" + echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT + else + echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT + fi + + - name: Release + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ steps.set_tag_name.outputs.tag_name }} + repository: sgl-project/whl + token: ${{ secrets.WHL_TOKEN }} + files: | + sgl-kernel/dist/* + + - name: Clone wheel index + run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl + env: + WHL_TOKEN: ${{ secrets.WHL_TOKEN }} + + - name: Update wheel index + run: python3 scripts/update_kernel_whl_index.py --cuda 128 + + - name: Push wheel index + run: | + cd sgl-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 index" + git push + + build-cu129-aarch64: + if: github.repository == 'sgl-project/sglang' + runs-on: sgl-kernel-release-node-arm + strategy: + matrix: + python-version: ["3.10"] + cuda-version: ["12.9"] + steps: + - uses: actions/checkout@v4 + with: + submodules: "recursive" + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Build wheels + run: | + cd sgl-kernel + chmod +x ./build.sh + ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" aarch64 + + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }}-aarch64 + path: sgl-kernel/dist/* + + release-cu129-aarch64: + needs: build-cu129-aarch64 + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + path: sgl-kernel/dist/ + merge-multiple: true + pattern: wheel-* + + - name: Set tag name + id: set_tag_name + run: | + if [ -z "${{ inputs.tag_name }}" ]; then + TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)" + echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT + else + echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT + fi + + - name: Release + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ steps.set_tag_name.outputs.tag_name }} + repository: sgl-project/whl + token: ${{ secrets.WHL_TOKEN }} + files: | + sgl-kernel/dist/* + + - name: Clone wheel index + run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl + env: + WHL_TOKEN: ${{ secrets.WHL_TOKEN }} + + - name: Update wheel index + run: python3 scripts/update_kernel_whl_index.py --cuda 129 + + - name: Push wheel index + run: | + cd sgl-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 index" + git push diff --git a/.github/workflows/vllm-dependency-test.yml b/.github/workflows/vllm-dependency-test.yml new file mode 100644 index 000000000..442d76d50 --- /dev/null +++ b/.github/workflows/vllm-dependency-test.yml @@ -0,0 +1,43 @@ +name: VLLM Dependency Test + +on: + push: + branches: [ main ] + paths: + - "python/**" + - "scripts/ci/**" + - "test/**" + pull_request: + branches: [ main ] + paths: + - "python/**" + - "scripts/ci/**" + - "test/**" + +concurrency: + group: vllm-dependency-test-${{ github.ref }} + cancel-in-progress: true + +jobs: + vllm-dependency-test: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + runs-on: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_dependency.sh + pip install "bitsandbytes>=0.44.0" + + pip install "sgl-kernel==0.3.7" + + - name: Run vLLM dependency tests + timeout-minutes: 60 + run: | + export SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1 + + cd test/srt + python3 run_suite.py --suite vllm_dependency_test --timeout-per-file 3600 diff --git a/.gitignore b/.gitignore new file mode 100644 index 000000000..9725fabd9 --- /dev/null +++ b/.gitignore @@ -0,0 +1,240 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ + +# Tokenizer cache for tests +.tokenizer_cache/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +.idea/ + +# MacOS +.DS_Store + +# Vim +*.swp + +# Documentation +docs/_build + +# SGL +benchmark/mmlu/data +benchmark/mmlu/data.tar +benchmark/llava_bench/images +benchmark/llava_bench/mme_pack +*.jsonl +tmp*.txt + +# Plots +*.png +*.pdf + +# personnal +work_dirs/ +*.csv + +!logo.png + +# Prerequisites +*.d + +# Compiled Object files +*.slo +*.lo +*.o +*.obj + +# Precompiled Headers +*.gch +*.pch + +# Compiled Dynamic libraries +*.so +*.dylib +*.dll + +# Fortran module files +*.mod +*.smod + +# Compiled Static libraries +*.lai +*.la +*.a +*.lib + +# Executables +*.exe +*.out +*.app + +compile_commands.json + +*.iml + +# VSCode +.vscode + +1 + +# Autoenv +.env.leave + +# Rust lib +Cargo.lock + +lmms-eval diff --git a/.isort.cfg b/.isort.cfg new file mode 100644 index 000000000..835095154 --- /dev/null +++ b/.isort.cfg @@ -0,0 +1,3 @@ +[settings] +profile=black +known_first_party=sglang diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 000000000..a295f2eb4 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,59 @@ +default_stages: [pre-commit, pre-push, manual] + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: check-symlinks + - id: destroyed-symlinks + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + args: [--allow-multiple-documents] + - id: check-toml + - id: check-ast + - id: check-added-large-files + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: detect-private-key + - id: debug-statements + - id: no-commit-to-branch + - repo: https://github.com/PyCQA/isort + rev: 5.13.2 + hooks: + - id: isort + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.11.7 + hooks: + - id: ruff + args: [--select=F401, --fixable=F401] + files: ^(benchmark/|docs/|examples/) + exclude: \.ipynb$ + - repo: https://github.com/psf/black + rev: 24.10.0 + hooks: + - id: black-jupyter + - repo: https://github.com/codespell-project/codespell + rev: v2.4.1 + hooks: + - id: codespell + additional_dependencies: ['tomli'] + args: ['--toml', 'python/pyproject.toml', '-L', 'cann,thi,makro,wil,rouge'] + exclude: | + (?x)^( + test/srt/test_reasoning_parser\.py| + docs/advanced_features/vlm_query\.ipynb + )$ + - repo: https://github.com/pre-commit/mirrors-clang-format + rev: v18.1.8 + hooks: + - id: clang-format + types_or: [c++, cuda] + args: [--style=file, --verbose] + - repo: https://github.com/kynan/nbstripout + rev: 0.8.1 + hooks: + - id: nbstripout + args: + - '--keep-output' + - '--extra-keys=metadata.kernelspec metadata.language_info.version' diff --git a/3rdparty/amd/profiling/PROFILING.md b/3rdparty/amd/profiling/PROFILING.md new file mode 100644 index 000000000..7e15ec844 --- /dev/null +++ b/3rdparty/amd/profiling/PROFILING.md @@ -0,0 +1,425 @@ +## Profiling SGLang Infer System with AMD GPUs +This AppNote describes the SGLang profiling technical, code augment and running steps for systems with AMD Instinct GPUs, nevertheless the same procedure may work with Nvidia GPUs too. +Examples and steps are provided in detail, to facilitate easy reproduce and use to localize performance problem towards optimizations. +Two primary methods are covered: +- [RPD](https://github.com/ROCm/rocmProfileData.git) +- [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) + +### Profiling SGLang Infer System with RPD Profiler +RPD profiler is a low-overhead cross-platform profiler. Therefore, the same RPD code augment not only works for profiling on ROCm/AMD GPUs, but also works for profiling on CUDA/Nvidia GPUs as well. To do RPD profiling on SGLang repository, please use scripts and patch files included in this directory and follow the steps below: +1. Install RPD with rpd.patch applied during installation using install_rpd.sh, both files are in this directory. + +install_rpd.sh + +```bash +# download and install RPD +apt update && apt install -y sqlite3 libsqlite3-dev libfmt-dev + +# install rpd module +git clone https://github.com/ROCmSoftwarePlatform/rocmProfileData +cd rocmProfileData +git checkout 976899e9c6dbc6dd2bccf770818e4e44125590ac +git apply rpd.patch +make && make install +cd rocpd_python && python setup.py install && cd .. +cd rpd_tracer && make clean;make install && python setup.py install && cd .. +``` + +rpd.patch + +```bash +diff --git a/rpd_tracer/Makefile b/rpd_tracer/Makefile +index e9d9feb..b2e9e1a 100644 +--- a/rpd_tracer/Makefile ++++ b/rpd_tracer/Makefile +@@ -16,7 +16,7 @@ ifneq (,$(HIP_PATH)) + $(info Building with roctracer) + RPD_LIBS += -L/opt/rocm/lib -lroctracer64 -lroctx64 -lamdhip64 -lrocm_smi64 + RPD_INCLUDES += -I/opt/rocm/include -I/opt/rocm/include/roctracer -I/opt/rocm/include/hsa +- RPD_SRCS += RoctracerDataSource.cpp RocmSmiDataSource.cpp ++ RPD_SRCS += RoctracerDataSource.cpp + RPD_INCLUDES += -D__HIP_PLATFORM_AMD__ + endif +``` +2. Add loadTracer.sh file included in this directory to /sglang/python/sglang. + +loadTracer.sh + +```bash +#!/bin/bash +################################################################################ +# Copyright (c) 2021 - 2023 Advanced Micro Devices, Inc. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +################################################################################ +OUTPUT_FILE="trace.rpd" + +if [ "$1" = "-o" ] ; then + OUTPUT_FILE=$2 + shift + shift +fi + +if [ -e ${OUTPUT_FILE} ] ; then + rm ${OUTPUT_FILE} +fi + +python3 -m rocpd.schema --create ${OUTPUT_FILE} +if [ $? != 0 ] ; then + echo "Error: Could not create rpd file. Please run 'python setup.py install' from the rocpd_python dir" + exit +fi + +export RPDT_FILENAME=${OUTPUT_FILE} +export RPDT_AUTOSTART=0 +LD_PRELOAD=librocm-smi_64:librpd_tracer.so "$@" +``` +3. Apply patch (provided in this directory) with "git apply rpd_profile_server_enable.patch" if the main profiling purpose is to get info on gpu kernels as well as limited cpu activity info. + +#### Common Notes 1 +Please note that although we are doing TP=8 in the example, we purposely only log RPD profiling on 2 ranks in the patch file (i.e.tp_rank=0/1) for profiling/visualization convenience, as even Perfetto streaming mode can only load maximal 8GB json file for visualization. With 2 ranks logged in RPD profiling, we could still check whether there are issues among ranks (e.g. load imbalance issue, nccl issue), and at the same time, we could log relatively longer time duration before the json file generated from RPD file hits 8GB size. + +rpd_profile_server_enable.patch + +```bash +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..9021c01 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -71,6 +71,8 @@ from sglang.srt.utils import ( + suppress_other_loggers, + ) + from sglang.utils import get_exception_traceback ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + logger = logging.getLogger(__name__) + +@@ -245,6 +247,7 @@ class Scheduler: + ], + with_stack=True, + ) ++ self.rpd = rpdTracerControl() + + @torch.inference_mode() + def event_loop(self): +@@ -1027,15 +1030,24 @@ class Scheduler: + def start_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.start() ++ #self.profiler.start() #block pytorch profiler for rpd profiler enabling ++ if self.tp_rank == 0 or self.tp_rank == 1: ++ self.rpd.start() ++ self.rpd.rangePush("", "rpd profile range", "") ++ logger.info("rpd is enabled") + + def stop_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ #self.profiler.stop() ++ #self.profiler.export_chrome_trace( ++ # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" ++ #) ++ if self.tp_rank ==0 or self.tp_rank ==1: ++ self.rpd.rangePop() ++ self.rpd.stop() ++ self.rpd.flush() ++ logger.info("rpd is done") + logger.info("Profiler is done") +``` + +#### Advanced Debugging with RPD Profiler +Sometimes, we want to use rpd profiler to capture more CPU and python activities in order to debug some challenging issues (e.g. root cause of load imbalance across gpu processes, root cause of bubbles, etc). Only in such cases, we need to apply patch "git apply rpd_profile_server_enable_wCPU_activities.patch", where 3 files are modified. + +rpd_profile_server_enable_wCPU_activities.patch + +```bash +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..2edb427 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -71,6 +71,8 @@ from sglang.srt.utils import ( + suppress_other_loggers, + ) + from sglang.utils import get_exception_traceback ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + logger = logging.getLogger(__name__) + +@@ -245,6 +247,7 @@ class Scheduler: + ], + with_stack=True, + ) ++ self.rpd = rpdTracerControl() + + @torch.inference_mode() + def event_loop(self): +@@ -1027,15 +1030,26 @@ class Scheduler: + def start_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.start() ++ #self.profiler.start() ++ logger.info("torch profiler is disabled") ++ if self.tp_rank == 0 or self.tp_rank == 1: ++ self.rpd.setPythonTrace(True) ++ self.rpd.start() ++ self.rpd.rangePush("", "scheduler", "") ++ logger.info("rpd is enabled inside scheduler profiling") + + def stop_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ #self.profiler.stop() ++ #self.profiler.export_chrome_trace( ++ # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" ++ #) ++ if self.tp_rank ==0 or self.tp_rank ==1: ++ self.rpd.rangePop() ++ self.rpd.stop() ++ self.rpd.flush() ++ logger.info("rpd is done inside scheduler") + logger.info("Profiler is done") + + +diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py +index 2621ccd..181df85 100644 +--- a/python/sglang/srt/managers/tokenizer_manager.py ++++ b/python/sglang/srt/managers/tokenizer_manager.py +@@ -58,6 +58,10 @@ from sglang.srt.sampling.sampling_params import SamplingParams + from sglang.srt.server_args import PortArgs, ServerArgs + from sglang.srt.utils import is_generation_model, is_multimodal_model + ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() ++ ++ + asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + + logger = logging.getLogger(__name__) +@@ -514,10 +518,20 @@ class TokenizerManager: + self.send_to_scheduler.send_pyobj(req) + + def start_profile(self): ++ rpd = rpdTracerControl() ++ rpd.setPythonTrace(True) ++ rpd.start() ++ rpd.rangePush("", "tokenizer_manager", "") ++ logger.info("tokenizer_manager rpd profiling started!") + req = ProfileReq.START_PROFILE + self.send_to_scheduler.send_pyobj(req) + + def stop_profile(self): ++ rpd = rpdTracerControl() ++ rpd.rangePop() ++ rpd.stop() ++ rpd.flush() ++ logger.info("rpd profiling is done inside tokenizer_manager!") + req = ProfileReq.STOP_PROFILE + self.send_to_scheduler.send_pyobj(req) + +diff --git a/python/sglang/srt/server.py b/python/sglang/srt/server.py +index 7111c93..2bd722c 100644 +--- a/python/sglang/srt/server.py ++++ b/python/sglang/srt/server.py +@@ -30,6 +30,8 @@ import threading + import time + from http import HTTPStatus + from typing import Dict, List, Optional, Union ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + # Fix a bug of Python threading + setattr(threading, "_register_atexit", lambda *args, **kwargs: None) +@@ -152,6 +154,11 @@ async def flush_cache(): + @app.post("/start_profile") + async def start_profile(): + """Start profiling.""" ++ rpd = rpdTracerControl() ++ rpd.setPythonTrace(True) ++ rpd.start() ++ rpd.rangePush("", "server rpd profile range", "") ++ logger.info("rpd profiling started in server.py!") + tokenizer_manager.start_profile() + return Response( + content="Start profiling.\n", +@@ -164,6 +171,11 @@ async def start_profile(): + async def stop_profile(): + """Stop profiling.""" + tokenizer_manager.stop_profile() ++ rpd = rpdTracerControl() ++ rpd.rangePop() ++ rpd.stop() ++ rpd.flush() ++ logger.info("rpd profiling is done in server.py!") + return Response( + content="Stop profiling. This will take some time.\n", + status_code=200, +``` + +4. As an example for grok1 profiling, we create a dummy_grok1 directory with config.json (see content below) inside this directory and copy this directory to the right path for "--model-path" if you want to use the example server.sh file provided. +```bash +cat ../dummy_grok1/config.json +{ + "architectures": [ + "Grok1ModelForCausalLM" + ], + "embedding_multiplier_scale": 78.38367176906169, + "output_multiplier_scale": 0.5773502691896257, + "vocab_size": 131072, + "hidden_size": 6144, + "intermediate_size": 32768, + "max_position_embeddings": 8192, + "num_experts_per_tok": 2, + "num_local_experts": 8, + "num_attention_heads": 48, + "num_hidden_layers": 64, + "num_key_value_heads": 8, + "head_dim": 128, + "rms_norm_eps": 1e-05, + "rope_theta": 10000.0, + "model_type": "mixtral", + "torch_dtype": "bfloat16" +} +``` +5. Launch server with rpd enabled script ./server.sh in one terminal inside the docker container. + +#### Common Notes 2 +- Remember to change model-path to the correct path +- loadTracer.sh is needed to conduct profiling +- SGLANG_TORCH_PROFILER_DIR is used for default torch profiler +- Do not use loadTracer.sh if you are using the torch profiler, simply use python3 -m sglang.launch_server. + + +server.sh + +```bash +#!/bin/bash + +# export SGLANG_TORCH_PROFILER_DIR=/data/sglang/ +export SGLANG_TORCH_PROFILER_DIR=/sgl-workspace/sglang/profile/ + +# Get the current timestamp +TIMESTAMP=$(date +"%Y%m%d_%H%M%S") + +# Define the log file with a timestamp +LOGFILE="sglang_server_log_$TIMESTAMP.json" + +# Run the Python command and save the output to the log file +loadTracer.sh python3 -m sglang.launch_server \ + --model-path /sgl-workspace/sglang/dummy_grok1 \ + --tokenizer-path Xenova/grok-1-tokenizer \ + --load-format dummy \ + --quantization fp8 \ + --tp 8 \ + --port 30000 \ + --disable-radix-cache 2>&1 | tee "$LOGFILE" +``` +6. Open another terminal for the same docker container, and run the rpd enabled ./client.sh after you see "The server is fired up and is ready to roll!" message from server side terminal. + +#### Common Notes 3 +- Use curl http://localhost:30000/start_profile & curl http://localhost:30000/stop_profile to control the start and end of profiling. Check sglang/python/sglang/srt/managers/scheduler.py for more details. +- Please don't use RPD profiler together with PyTorch profiler to avoid interference. +- The rocmProfileData/tools/rpd2tracing.py file is used to generate json file from RPD file. + +client.sh + +```bash +#!/bin/bash + +# Start profiling via API +curl http://localhost:30000/start_profile -H "Content-Type: application/json" + +# Benchmark serving using sglang with random dataset and tokenizer +# Define the log file with a timestamp +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +LOGFILE="sglang_client_log_$TIMESTAMP.json" + +# Run the benchmark with specified parameters and save logs +python3 -m sglang.bench_serving \ + --backend sglang \ + --tokenizer Xenova/grok-1-tokenizer \ + --dataset-name random \ + --random-input 1024\ + --random-output 1024 \ + --num-prompts 120 \ + --request-rate 8 \ + --output-file online.jsonl 2>&1 | tee "$LOGFILE" + +# Stop profiling via API +curl http://localhost:30000/stop_profile -H "Content-Type: application/json" + +# Convert tracing file to csv & json +sqlite3 trace.rpd ".mode csv" ".header on" ".output trace.csv" "select * from top;" ".output stdout" +python3 ./rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json +``` +7. Follow [Perfetto docs](https://perfetto.dev/docs/visualization/large-traces) to visualize large json files. Try to adjust parameters so that the trace.json file size is less than 9GB. + +### Profiling SGLang Infer System with PyTorch Profiler + +Please use the steps as follows: + +1. Apply the patch torch_profiler.patch. Note that you can modify "if self.tp_rank == 0" in the patch to allow more ranks be recorded in profiling. + +torch_profiler.patch +```bash +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..6ecd78c 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -240,7 +240,6 @@ class Scheduler: + ) + self.profiler = torch.profiler.profile( + activities=[ +- torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, +@@ -1033,9 +1032,11 @@ class Scheduler: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") + self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ if self.tp_rank == 0: ++ with open(f"stats_repro_{int(time.time())}.txt", "w") as f: ++ print(self.profiler.key_averages(group_by_input_shape=True).table(sort_by="cuda_time_total", row_limit=-1), file=f) ++ print("Profiling stats done.") ++ + logger.info("Profiler is done") +``` + +2. Create the model path directory and copy it to the right path for "--model-path" if you want to use the server.sh file provided. + +3. Modify the included server.sh by removing "loadTracer.sh" before python command and launch script ./server.sh in one terminal inside the docker container. + +4. Similar to step 6 in RPD profiling section, but remove the last 2 lines in client.sh, which converted rpd file into csv and json files. Run modified client.sh for PyTorch profiling. +------- +- [Torch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) diff --git a/3rdparty/amd/profiling/client.sh b/3rdparty/amd/profiling/client.sh new file mode 100755 index 000000000..150ea9f19 --- /dev/null +++ b/3rdparty/amd/profiling/client.sh @@ -0,0 +1,27 @@ +#!/bin/bash + +# Start profiling via API +curl http://localhost:30000/start_profile -H "Content-Type: application/json" + +# Benchmark serving using sglang with random dataset and tokenizer +# Define the log file with a timestamp +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +LOGFILE="sglang_client_log_$TIMESTAMP.json" + +# Run the benchmark with specified parameters and save logs +python3 -m sglang.bench_serving \ + --backend sglang \ + --tokenizer Xenova/grok-1-tokenizer \ + --dataset-name random \ + --random-input 1024\ + --random-output 1024 \ + --num-prompts 240 \ + --request-rate 8 \ + --output-file online.jsonl 2>&1 | tee "$LOGFILE" + +# Stop profiling via API +curl http://localhost:30000/stop_profile -H "Content-Type: application/json" + +# Convert tracing file to csv & json +sqlite3 trace.rpd ".mode csv" ".header on" ".output trace.csv" "select * from top;" ".output stdout" +python3 /sgl-workspace/rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json diff --git a/3rdparty/amd/profiling/install_rpd.sh b/3rdparty/amd/profiling/install_rpd.sh new file mode 100644 index 000000000..d1b04b988 --- /dev/null +++ b/3rdparty/amd/profiling/install_rpd.sh @@ -0,0 +1,10 @@ +# download and install RPD +apt update && apt install -y sqlite3 libsqlite3-dev libfmt-dev + +# install rpd module +git clone https://github.com/ROCmSoftwarePlatform/rocmProfileData +cd rocmProfileData +git apply rpd.patch +make && make install +cd rocpd_python && python setup.py install && cd .. +cd rpd_tracer && make clean;make install && python setup.py install && cd .. diff --git a/3rdparty/amd/profiling/loadTracer.sh b/3rdparty/amd/profiling/loadTracer.sh new file mode 100755 index 000000000..8a95a335c --- /dev/null +++ b/3rdparty/amd/profiling/loadTracer.sh @@ -0,0 +1,43 @@ +#!/bin/bash +################################################################################ +# Copyright (c) 2021 - 2023 Advanced Micro Devices, Inc. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +################################################################################ +OUTPUT_FILE="trace.rpd" + +if [ "$1" = "-o" ] ; then + OUTPUT_FILE=$2 + shift + shift +fi + +if [ -e ${OUTPUT_FILE} ] ; then + rm ${OUTPUT_FILE} +fi + +python3 -m rocpd.schema --create ${OUTPUT_FILE} +if [ $? != 0 ] ; then + echo "Error: Could not create rpd file. Please run 'python setup.py install' from the rocpd_python dir" + exit +fi + +export RPDT_FILENAME=${OUTPUT_FILE} +export RPDT_AUTOSTART=0 +LD_PRELOAD=librocm-smi_64:librpd_tracer.so "$@" diff --git a/3rdparty/amd/profiling/rpd.patch b/3rdparty/amd/profiling/rpd.patch new file mode 100644 index 000000000..87917654a --- /dev/null +++ b/3rdparty/amd/profiling/rpd.patch @@ -0,0 +1,12 @@ +diff --git a/rpd_tracer/Makefile b/rpd_tracer/Makefile +index e9d9feb..b2e9e1a 100644 +--- a/rpd_tracer/Makefile ++++ b/rpd_tracer/Makefile +@@ -16,7 +16,7 @@ ifneq (,$(HIP_PATH)) + $(info Building with roctracer) + RPD_LIBS += -L/opt/rocm/lib -lroctracer64 -lroctx64 -lamdhip64 -lrocm_smi64 + RPD_INCLUDES += -I/opt/rocm/include -I/opt/rocm/include/roctracer -I/opt/rocm/include/hsa +- RPD_SRCS += RoctracerDataSource.cpp RocmSmiDataSource.cpp ++ RPD_SRCS += RoctracerDataSource.cpp + RPD_INCLUDES += -D__HIP_PLATFORM_AMD__ + endif diff --git a/3rdparty/amd/profiling/rpd_profile_server_enable.patch b/3rdparty/amd/profiling/rpd_profile_server_enable.patch new file mode 100644 index 000000000..3cd391533 --- /dev/null +++ b/3rdparty/amd/profiling/rpd_profile_server_enable.patch @@ -0,0 +1,49 @@ +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..9021c01 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -71,6 +71,8 @@ from sglang.srt.utils import ( + suppress_other_loggers, + ) + from sglang.utils import get_exception_traceback ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + logger = logging.getLogger(__name__) + +@@ -245,6 +247,7 @@ class Scheduler: + ], + with_stack=True, + ) ++ self.rpd = rpdTracerControl() + + @torch.inference_mode() + def event_loop(self): +@@ -1027,15 +1030,24 @@ class Scheduler: + def start_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.start() ++ #self.profiler.start() #block pytorch profiler for rpd profiler enabling ++ if self.tp_rank == 0 or self.tp_rank == 1: ++ self.rpd.start() ++ self.rpd.rangePush("", "rpd profile range", "") ++ logger.info("rpd is enabled") + + def stop_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ #self.profiler.stop() ++ #self.profiler.export_chrome_trace( ++ # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" ++ #) ++ if self.tp_rank ==0 or self.tp_rank ==1: ++ self.rpd.rangePop() ++ self.rpd.stop() ++ self.rpd.flush() ++ logger.info("rpd is done") + logger.info("Profiler is done") diff --git a/3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch b/3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch new file mode 100644 index 000000000..5416f4d57 --- /dev/null +++ b/3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch @@ -0,0 +1,126 @@ +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..2edb427 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -71,6 +71,8 @@ from sglang.srt.utils import ( + suppress_other_loggers, + ) + from sglang.utils import get_exception_traceback ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + logger = logging.getLogger(__name__) + +@@ -245,6 +247,7 @@ class Scheduler: + ], + with_stack=True, + ) ++ self.rpd = rpdTracerControl() + + @torch.inference_mode() + def event_loop(self): +@@ -1027,15 +1030,26 @@ class Scheduler: + def start_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.start() ++ #self.profiler.start() ++ logger.info("torch profiler is disabled") ++ if self.tp_rank == 0 or self.tp_rank == 1: ++ self.rpd.setPythonTrace(True) ++ self.rpd.start() ++ self.rpd.rangePush("", "scheduler", "") ++ logger.info("rpd is enabled inside scheduler profiling") + + def stop_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ #self.profiler.stop() ++ #self.profiler.export_chrome_trace( ++ # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" ++ #) ++ if self.tp_rank ==0 or self.tp_rank ==1: ++ self.rpd.rangePop() ++ self.rpd.stop() ++ self.rpd.flush() ++ logger.info("rpd is done inside scheduler") + logger.info("Profiler is done") + + +diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py +index 2621ccd..181df85 100644 +--- a/python/sglang/srt/managers/tokenizer_manager.py ++++ b/python/sglang/srt/managers/tokenizer_manager.py +@@ -58,6 +58,10 @@ from sglang.srt.sampling.sampling_params import SamplingParams + from sglang.srt.server_args import PortArgs, ServerArgs + from sglang.srt.utils import is_generation_model, is_multimodal_model + ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() ++ ++ + asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + + logger = logging.getLogger(__name__) +@@ -514,10 +518,20 @@ class TokenizerManager: + self.send_to_scheduler.send_pyobj(req) + + def start_profile(self): ++ rpd = rpdTracerControl() ++ rpd.setPythonTrace(True) ++ rpd.start() ++ rpd.rangePush("", "tokenizer_manager", "") ++ logger.info("tokenizer_manager rpd profiling started!") + req = ProfileReq.START_PROFILE + self.send_to_scheduler.send_pyobj(req) + + def stop_profile(self): ++ rpd = rpdTracerControl() ++ rpd.rangePop() ++ rpd.stop() ++ rpd.flush() ++ logger.info("rpd profiling is done inside tokenizer_manager!") + req = ProfileReq.STOP_PROFILE + self.send_to_scheduler.send_pyobj(req) + +diff --git a/python/sglang/srt/server.py b/python/sglang/srt/server.py +index 7111c93..2bd722c 100644 +--- a/python/sglang/srt/server.py ++++ b/python/sglang/srt/server.py +@@ -30,6 +30,8 @@ import threading + import time + from http import HTTPStatus + from typing import Dict, List, Optional, Union ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + # Fix a bug of Python threading + setattr(threading, "_register_atexit", lambda *args, **kwargs: None) +@@ -152,6 +154,11 @@ async def flush_cache(): + @app.post("/start_profile") + async def start_profile(): + """Start profiling.""" ++ rpd = rpdTracerControl() ++ rpd.setPythonTrace(True) ++ rpd.start() ++ rpd.rangePush("", "server rpd profile range", "") ++ logger.info("rpd profiling started in server.py!") + tokenizer_manager.start_profile() + return Response( + content="Start profiling.\n", +@@ -164,6 +171,11 @@ async def start_profile(): + async def stop_profile(): + """Stop profiling.""" + tokenizer_manager.stop_profile() ++ rpd = rpdTracerControl() ++ rpd.rangePop() ++ rpd.stop() ++ rpd.flush() ++ logger.info("rpd profiling is done in server.py!") + return Response( + content="Stop profiling. This will take some time.\n", + status_code=200, diff --git a/3rdparty/amd/profiling/server.sh b/3rdparty/amd/profiling/server.sh new file mode 100755 index 000000000..f877e6c7a --- /dev/null +++ b/3rdparty/amd/profiling/server.sh @@ -0,0 +1,20 @@ +#!/bin/bash + +# export SGLANG_TORCH_PROFILER_DIR=/data/sglang/ +export SGLANG_TORCH_PROFILER_DIR=/sgl-workspace/sglang/profile/ + +# Get the current timestamp +TIMESTAMP=$(date +"%Y%m%d_%H%M%S") + +# Define the log file with a timestamp +LOGFILE="sglang_server_log_$TIMESTAMP.json" + +# Run the Python command and save the output to the log file +loadTracer.sh python3 -m sglang.launch_server \ + --model-path /sgl-workspace/sglang/dummy_grok1 \ + --tokenizer-path Xenova/grok-1-tokenizer \ + --load-format dummy \ + --quantization fp8 \ + --tp 8 \ + --port 30000 \ + --disable-radix-cache 2>&1 | tee "$LOGFILE" diff --git a/3rdparty/amd/profiling/torch_profiler.patch b/3rdparty/amd/profiling/torch_profiler.patch new file mode 100644 index 000000000..40f557406 --- /dev/null +++ b/3rdparty/amd/profiling/torch_profiler.patch @@ -0,0 +1,25 @@ +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..6ecd78c 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -240,7 +240,6 @@ class Scheduler: + ) + self.profiler = torch.profiler.profile( + activities=[ +- torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, +@@ -1033,9 +1032,11 @@ class Scheduler: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") + self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ if self.tp_rank == 0: ++ with open(f"stats_repro_{int(time.time())}.txt", "w") as f: ++ print(self.profiler.key_averages(group_by_input_shape=True).table(sort_by="cuda_time_total", row_limit=-1), file=f) ++ print("Profiling stats done.") ++ + logger.info("Profiler is done") diff --git a/3rdparty/amd/tuning/TUNING.md b/3rdparty/amd/tuning/TUNING.md new file mode 100644 index 000000000..e7b9b2049 --- /dev/null +++ b/3rdparty/amd/tuning/TUNING.md @@ -0,0 +1,118 @@ +## Tuning SGLang Infer System with AMD GPUs +This AppNote describes the SGLang performance tuning technical, code harness and running steps for systems with AMD Instinct GPUs. +Harness code, examples and steps are provided in detail, to facilitate easy reproduce & use to tune performance towards workloads. +Three primary runtime areas are covered: + +## 1. Triton Kernels +To maximize Triton kernel efficiency, several strategies can be employed: + +### Key Environment Variables: +- **num_stages**: Adjusts the number of pipeline stages to optimize kernel efficiency based on the specific type of operations (e.g., General Matrix Multiplication - GEMM). +- **waves_per_eu**: Controls the usage of Vector General Purpose Registers (VGPR) to enhance occupancy, thereby improving latency or throughput. +- **BLOCK_M, BLOCK_N, BLOCK_K**: Tunable tile sizes that assist in balancing memory transfer and computational efficiency. +- **matrix_instr_nonkdim**: Optimizes the usage of Matrix-Fused Multiply-Add (MFMA) instructions for specific kernel types, such as Flash Attention. +- **OPTIMIZE_EPILOGUE**: An environment variable that can be set to `1` to enhance performance by eliminating the `convert_layout` operation in the kernel's epilogue. +```python +@triton.autotune(configs=[ + triton.Config({'waves_per_eu': 1}, num_warps=4, num_stages=1), + triton.Config({'waves_per_eu': 1}, num_warps=8, num_stages=1), + triton.Config({'waves_per_eu': 1}, num_warps=16, num_stages=1), + triton.Config({'waves_per_eu': 2}, num_warps=4, num_stages=1), + triton.Config({'waves_per_eu': 2}, num_warps=8, num_stages=1), + triton.Config({'waves_per_eu': 2}, num_warps=16, num_stages=1), + triton.Config({'waves_per_eu': 4}, num_warps=4, num_stages=1), + triton.Config({'waves_per_eu': 4}, num_warps=8, num_stages=1), + triton.Config({'waves_per_eu': 4}, num_warps=16, num_stages=1), + ], key=['BLOCK_N', 'NUM_TOKEN_BLKS'], use_cuda_graph=True) +@triton.jit +def _triton_kernel_funtion(): + ... +``` +## 2. Torch Tunable Operations +**TunableOp** is a feature in PyTorch that allows for the definition and optimization of custom kernels with tunable parameters. This feature is particularly useful for enhancing the performance of kernels by experimenting with different configurations. + +### Key Environment Variables: +1. **PYTORCH_TUNABLEOP_ENABLED**: + - Default: `0` + - Set to `1` to enable TunableOp. + +2. **PYTORCH_TUNABLEOP_TUNING**: + - Default: `1` + - Set to `0` to disable tuning. If a tuned entry is not found, it will run the tuning step and record the entry when PYTORCH_TUNABLEOP_ENABLED is enabled. + +3. **PYTORCH_TUNABLEOP_VERBOSE**: + - Default: `0` + - Set to `1` to enable verbose output for TunableOp. + +### Usage Example: +To enable TunableOp and tuning, and optionally enable verbose mode, you can run the following command in your terminal: + +```bash +#Tuning +PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=1 your_script.sh + +#Inference with tuning op +PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=0 your_script.sh + +#Print out the log +PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=0 PYTORCH_TUNABLEOP_VERBOSE=1 your_script.sh + +``` +## 3. Torch Compilation + + +The following are suggestions for optimizing matrix multiplication (GEMM) and convolution (conv) operations in PyTorch using Inductor, a part of the PyTorch compilation framework. The goal is to leverage Triton to achieve better performance. + +To tune Triton kernels with GEMM and convolution ops (conv), use the `torch.compile` function with the max-autotune mode. This benchmarks a predefined list of Triton configurations and selects the fastest one for each shape. + +### Key Configurations: +1. **Max Autotune**: + - Set `torch._inductor.config.max_autotune = True` or `TORCHINDUCTOR_MAX_AUTOTUNE=1`. + +2. **Fine-Grained Control**: + - Enable GEMM tuning: `torch._inductor.config.max_autotune_gemm = True`. + - Enable tuning for pointwise/reduction ops: `torch._inductor.config.max_autotune.pointwise = True`. + +3. **Backend Selection**: + - Use `torch._inductor.max_autotune_gemm_backends` to limit backends to TRITON for better performance. + +4. **Freezing for Inference**: + - Use `torch._inductor.config.freezing=True` to enable constant folding optimizations. + +5. **Debugging**: + - Set `TORCH_COMPILE_DEBUG=1` to extract Triton kernels generated by Inductor. + +### Example Code Block: +```bash +#Gemm Tuning +TORCHINDUCTOR_MAX_AUTOTUNE=1 TORCHINDUCTOR_COORDINATE_DESCENT_TUNING=1 your_script.sh + +#Specify your backend to TRITON for Gemm Tuning +TORCHINDUCTOR_MAX_AUTOTUNE=1 TORCHINDUCTOR_COORDINATE_DESCENT_TUNING=1 TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS=TRITON your_script.sh + +#Inference with large improvement on AMD GPU +TORCHINDUCTOR_FREEZING=1 your_script.sh +``` +## 4. Fused MOE kernel +To maximize moe kernel efficiency, need to use below scripts to find out the best launch configuration + +### Key parameters: +- **--model**: what moe model type to do tuning, it will automatically decide the size of d_model, model_intermediate_size, num_layers +- **--tp-size**: simulate the whole model run configuration to set the dimension size using tp correctly +- **--batch**: M dimension size of moe kernel, for prefill moe kernel the value is batch*input_len, for decode moe kernel the value is batch +- **--dtype**: computation type + +```bash +#Tuning +#for example, we have one case like this "python3 -m sglang.bench_latency --model dummy_grok1/ --load-format dummy --tokenizer-path Xenova/grok-1-tokenizer --tp 8 --batch-size 32 --input 1024 --output 8 --attention-backend triton --sampling-backend pytorch --quantization fp8" to run, it defined batch-size 32 input length 1024 and output length 8, from "--batch" in moe view point, the prefill batch is 32*1024 = 32768, the decode batch is 32*1(only one output token generated in each run). +#so we can tune decode moe use below command +python benchmark_moe_rocm.py --model grok1 --tp-size 8 --dtype float8 --batch "32" +# and use this command to tune prefill moe +python benchmark_moe_rocm.py --model grok1 --tp-size 8 --dtype float8 --batch "32768" +``` + +## Reference + +For more detailed information on tuning SGLang performance with AMD GPUs, please refer to the following link: + +[ROCm Documentation: Triton Kernel Performance Optimization](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html#triton-kernel-performance-optimization) diff --git a/3rdparty/amd/tuning/benchmark_moe_rocm.py b/3rdparty/amd/tuning/benchmark_moe_rocm.py new file mode 100644 index 000000000..af596d218 --- /dev/null +++ b/3rdparty/amd/tuning/benchmark_moe_rocm.py @@ -0,0 +1,380 @@ +import argparse +import json +import os +import sys + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from tqdm import tqdm +from transformers import AutoConfig + +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + fused_moe, + get_config_file_name, +) + +padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0 + + +def main(model, tp_size, dtype: str, batches): + method = fused_moe + + for bs in batches: + run_grid(int(bs), model=model, method=method, tp_size=tp_size, dtype=dtype) + + +def prune_configs(M, N, K, configs): + pruned_configs = [] + elemBytes_a = 1 # [DV Note] Hard-coded for float16 (2 bytes) + elemBytes_b = 1 # [DV Note] Hard-coded for float16 (2 bytes) + + mfma = 16 if M < 32 or N < 32 else 32 + + # TODO (zhanglx): figure out the boundary between large and small gemms + large_gemm = False + if M >= 2048 and N >= 2048: + large_gemm = True + + for config in configs: + BLOCK_SIZE_M = config.get("BLOCK_SIZE_M") + BLOCK_SIZE_N = config.get("BLOCK_SIZE_N") + BLOCK_SIZE_K = config.get("BLOCK_SIZE_K") + num_warps = config.get("num_warps") + matrix_instr_nonkdim = config.get("matrix_instr_nonkdim") + # kpack = config.get("kpack") + if matrix_instr_nonkdim > mfma: + continue + if mfma == 4 and BLOCK_SIZE_K < 64: + continue + # some layouts could not work properly in case + # number elements per thread is less 1 + if BLOCK_SIZE_M * BLOCK_SIZE_N < 64: + continue + SPLIT_K = 1 # config.get("SPLIT_K") + GROUP_M = config.get("GROUP_SIZE_M") + if matrix_instr_nonkdim > BLOCK_SIZE_M or matrix_instr_nonkdim > BLOCK_SIZE_N: + continue + if matrix_instr_nonkdim >= M and matrix_instr_nonkdim != BLOCK_SIZE_M: + continue + if matrix_instr_nonkdim >= N and matrix_instr_nonkdim != BLOCK_SIZE_N: + continue + # Skip BLOCK_SIZE that is too large compare to M/N + # unless BLOCK_SIZE is already small enough + if M * 2 < BLOCK_SIZE_M and BLOCK_SIZE_M != 16: + continue + if N * 2 < BLOCK_SIZE_N and BLOCK_SIZE_N != 16: + continue + # skip large split_k when not necessary + if SPLIT_K != 1 and not need_split_k(M, N, K): + continue + # skip split_k that leads to EVEN_K = false + leap = SPLIT_K * BLOCK_SIZE_K + modv = K % leap + if modv != 0: + continue + # skip large GROUP_M + if GROUP_M * BLOCK_SIZE_M > M and GROUP_M != 1: + continue + # out of shared memory resource + # TODO (zhanglx): This does not consider the LDS usage in the epilogue + LDS = ( + BLOCK_SIZE_K * BLOCK_SIZE_M * elemBytes_a + + BLOCK_SIZE_K * BLOCK_SIZE_N * elemBytes_b + ) + if LDS > 65536: + continue + # Skip small block sizes and num_warps for large gemm + # For fp16 and f8, we want to only use BLOCK_SIZE >= 64 + if large_gemm: + if BLOCK_SIZE_M < 64 or BLOCK_SIZE_N < 64: + continue + if BLOCK_SIZE_K < 64: + continue + if num_warps < 4: + continue + + pruned_configs.append(config) + + return pruned_configs + + +def union_of_list_of_dicts(l1, l2): + result = [] + temp_list = l1.copy() + temp_list.extend(l2) + for myDict in temp_list: + if myDict not in result: + result.append(myDict) + + return result + + +def run_grid(bs, model, method, tp_size, dtype: str): + + config = AutoConfig.from_pretrained(model) + + top_k = config.num_experts_per_tok + d_model = config.hidden_size + model_intermediate_size = config.intermediate_size + num_layers = config.num_hidden_layers + hidden_states_dtype = config.torch_dtype + + if config.num_experts_per_tok: + if config.architectures[0] == "Grok1ModelForCausalLM": + num_total_experts = config.num_experts + else: + num_total_experts = config.num_local_experts + else: + raise ValueError(f"Unsupported Mixtral model {model}") + + # tp_size = 2 + num_warmup_calls = 10 + num_calls = 30 + + num_warmup_trials = 1 + num_trials = 1 + + full_configs = [] + + block_m_range = [16, 32, 64, 128, 256] + block_n_range = [16, 32, 64, 128, 256] + block_k_range = [32, 64, 128, 256] # MUST >= 32 + num_warps_range = [1, 2, 4, 8] + group_m_range = [1, 4, 8, 16, 32] + # For now we see better perf with num_stages=0 for all gemm configs we care + # But keep this explicit so that we do not forget we may need to set it to + # other values in the future + num_stage_range = [2] + waves_per_eu_range = [0, 1, 2, 4, 8] + # Remove 32 because of triton compiling error + matrix_instr_nonkdim_range = [16] + kpack_range = [1, 2] + + for block_size_m in block_m_range: + for block_size_n in block_n_range: + for block_size_k in block_k_range: + for group_size_m in group_m_range: + for num_warps in num_warps_range: + for num_stages in num_stage_range: + for waves_per_eu in waves_per_eu_range: + for matrix_instr_nonkdim in matrix_instr_nonkdim_range: + for kpack in kpack_range: + full_configs.append( + { + "BLOCK_SIZE_M": block_size_m, + "BLOCK_SIZE_N": block_size_n, + "BLOCK_SIZE_K": block_size_k, + "GROUP_SIZE_M": group_size_m, + "num_warps": num_warps, + "num_stages": num_stages, + "waves_per_eu": waves_per_eu, + "matrix_instr_nonkdim": matrix_instr_nonkdim, + "kpack": kpack, + } + ) + + M1 = bs * 2 + N1 = model_intermediate_size * 2 // tp_size + K1 = d_model + prune_configs_1 = prune_configs(M1, N1, K1, full_configs) + + M2 = bs * 2 + N2 = d_model + K2 = model_intermediate_size // tp_size + prune_configs_2 = prune_configs(M2, N2, K2, full_configs) + + configs = union_of_list_of_dicts(prune_configs_1, prune_configs_2) + + print( + f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \ + {len(prune_configs_2)=} | {len(configs)=}" + ) + + best_config = None + best_time_us = 1e20 + + print(f"{tp_size=} {bs=}") + + for config in tqdm(configs): + # warmup + try: + print(config) + for _ in range(num_warmup_trials): + run_timing( + num_calls=num_warmup_calls, + bs=bs, + d_model=d_model, + num_total_experts=num_total_experts, + top_k=top_k, + tp_size=tp_size, + model_intermediate_size=model_intermediate_size, + method=method, + config=config, + dtype=dtype, + hidden_states_dtype=hidden_states_dtype, + ) + except triton.runtime.autotuner.OutOfResources: + continue + + # trial + for _ in range(num_trials): + kernel_dur_ms = run_timing( + num_calls=num_calls, + bs=bs, + d_model=d_model, + num_total_experts=num_total_experts, + top_k=top_k, + tp_size=tp_size, + model_intermediate_size=model_intermediate_size, + method=method, + config=config, + dtype=dtype, + hidden_states_dtype=hidden_states_dtype, + ) + + kernel_dur_us = 1000 * kernel_dur_ms + model_dur_ms = kernel_dur_ms * num_layers + + if kernel_dur_us < best_time_us: + best_config = config + best_time_us = kernel_dur_us + + tqdm.write( + f"{kernel_dur_us=:.1f} {model_dur_ms=:.1f}" + f" {bs=} {tp_size=} {top_k=} {num_total_experts=} " + f"{d_model=} {model_intermediate_size=} {num_layers=}" + ) + + print("best_time_us", best_time_us) + print("best_config", best_config) + + # holds Dict[str, Dict[str, int]] + filename = get_config_file_name( + num_total_experts, + model_intermediate_size // tp_size, + "float8" if dtype == "float8" else None, + ) + print(f"writing config to file {filename}") + existing_content = {} + if os.path.exists(filename): + with open(filename, "r") as f: + existing_content = json.load(f) + existing_content[str(bs)] = best_config + with open(filename, "w") as f: + json.dump(existing_content, f, indent=4) + f.write("\n") + + +def run_timing( + num_calls: int, + bs: int, + d_model: int, + num_total_experts: int, + top_k: int, + tp_size: int, + model_intermediate_size: int, + method, + config, + dtype: str, + hidden_states_dtype, +) -> float: + shard_intermediate_size = model_intermediate_size // tp_size + + hidden_states = torch.rand( + (bs, d_model), + device="cuda:0", + dtype=hidden_states_dtype, + ) + + w1 = torch.rand( + (num_total_experts, 2 * shard_intermediate_size, d_model + padding_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + w2 = torch.rand( + (num_total_experts, d_model, shard_intermediate_size + padding_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + w1_scale = None + w2_scale = None + a1_scale = None + a2_scale = None + + if dtype == "float8": + w1 = w1.to(torch.float8_e4m3fnuz) + w2 = w2.to(torch.float8_e4m3fnuz) + w1_scale = torch.ones( + num_total_experts, device=hidden_states.device, dtype=torch.float32 + ) + w2_scale = torch.ones( + num_total_experts, device=hidden_states.device, dtype=torch.float32 + ) + a1_scale = torch.ones(1, device=hidden_states.device, dtype=torch.float32) + a2_scale = torch.ones(1, device=hidden_states.device, dtype=torch.float32) + + gating_output = F.softmax( + torch.rand( + (num_calls, bs, num_total_experts), + device=hidden_states.device, + dtype=torch.float32, + ), + dim=-1, + ) + + ################################## + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + start_event.record() + for i in range(num_calls): + hidden_states = method( + hidden_states=hidden_states, + w1=w1, + w2=w2, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + gating_output=gating_output[0], + topk=top_k, + renormalize=True, + inplace=True, + override_config=config, + use_fp8=dtype == "float8", + ) + + end_event.record() + end_event.synchronize() + + dur_ms = start_event.elapsed_time(end_event) / num_calls + return dur_ms + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="benchmark_mixtral_moe", + description="Benchmark and tune the fused_moe kernel", + ) + parser.add_argument( + "--dtype", + type=str, + default="auto", + choices=["float8", "float16", "bfloat16"], + help="Data type used for fused_moe kernel computations", + ) + parser.add_argument("--model", type=str, default="hpcai-tech/grok-1") + + parser.add_argument("--tp-size", type=int, default=2, help="Tensor paralleli size") + parser.add_argument("-b", "--batches", type=str) + + args = parser.parse_args() + + batches = args.batches.split(",") + + sys.exit(main(args.model, args.tp_size, args.dtype, batches)) diff --git a/LICENSE b/LICENSE new file mode 100644 index 000000000..9c422689c --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2023-2024 SGLang 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/Makefile b/Makefile new file mode 100644 index 000000000..459dfa573 --- /dev/null +++ b/Makefile @@ -0,0 +1,46 @@ +.PHONY: check-deps install-deps format update help + +# Show help for each target +help: + @echo "Available targets:" + @grep -E '^[a-zA-Z0-9_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}' + +check-deps: ## Check and install required Python formatting dependencies + @command -v isort >/dev/null 2>&1 || (echo "Installing isort..." && pip install isort) + @command -v black >/dev/null 2>&1 || (echo "Installing black..." && pip install black) + +install-deps: ## Install Python formatting tools (isort and black) + pip install isort black + +format: check-deps ## Format modified Python files using isort and black + @echo "Formatting modified Python files..." + git diff --name-only --diff-filter=M | grep '\.py$$' | xargs -I {} sh -c 'isort {} && black {}' + +FILES_TO_UPDATE = docker/Dockerfile.rocm \ + python/pyproject.toml \ + python/sglang/version.py \ + docs/developer_guide/setup_github_runner.md \ + docs/get_started/install.md \ + docs/platforms/amd_gpu.md \ + docs/platforms/ascend_npu.md \ + benchmark/deepseek_v3/README.md + +update: ## Update version numbers across project files. Usage: make update + @if [ -z "$(filter-out $@,$(MAKECMDGOALS))" ]; then \ + echo "Version required. Usage: make update "; \ + exit 1; \ + fi + @OLD_VERSION=$$(grep "version" python/sglang/version.py | cut -d '"' -f2); \ + NEW_VERSION=$(filter-out $@,$(MAKECMDGOALS)); \ + echo "Updating version from $$OLD_VERSION to $$NEW_VERSION"; \ + for file in $(FILES_TO_UPDATE); do \ + if [ "$(shell uname)" = "Darwin" ]; then \ + sed -i '' -e "s/$$OLD_VERSION/$$NEW_VERSION/g" $$file; \ + else \ + sed -i -e "s/$$OLD_VERSION/$$NEW_VERSION/g" $$file; \ + fi \ + done; \ + echo "Version update complete" + +%: + @: diff --git a/README.md b/README.md new file mode 100644 index 000000000..451a6d424 --- /dev/null +++ b/README.md @@ -0,0 +1,78 @@ +
+logo + +[![PyPI](https://img.shields.io/pypi/v/sglang)](https://pypi.org/project/sglang) +![PyPI - Downloads](https://static.pepy.tech/badge/sglang?period=month) +[![license](https://img.shields.io/github/license/sgl-project/sglang.svg)](https://github.com/sgl-project/sglang/tree/main/LICENSE) +[![issue resolution](https://img.shields.io/github/issues-closed-raw/sgl-project/sglang)](https://github.com/sgl-project/sglang/issues) +[![open issues](https://img.shields.io/github/issues-raw/sgl-project/sglang)](https://github.com/sgl-project/sglang/issues) +[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/sgl-project/sglang) + +
+ +-------------------------------------------------------------------------------- + +| [**Blog**](https://lmsys.org/blog/2025-05-05-large-scale-ep/) +| [**Documentation**](https://docs.sglang.ai/) +| [**Join Slack**](https://slack.sglang.ai/) +| [**Join Bi-Weekly Development Meeting**](https://meeting.sglang.ai/) +| [**Roadmap**](https://github.com/sgl-project/sglang/issues/7736) +| [**Slides**](https://github.com/sgl-project/sgl-learning-materials?tab=readme-ov-file#slides) | + +## News +- [2025/08] 🔔 SGLang x AMD SF Meetup on 8/22: Hands-on GPU workshop, tech talks by AMD/xAI/SGLang, and networking ([Roadmap](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_sglang_roadmap.pdf), [Large-scale EP](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_sglang_ep.pdf), [Highlights](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_highlights.pdf), [AITER/MoRI](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_aiter_mori.pdf), [Wave](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_wave.pdf)). +- [2025/08] 🔥 SGLang provides day-0 support for OpenAI gpt-oss model ([instructions](https://github.com/sgl-project/sglang/issues/8833)) +- [2025/06] 🔥 SGLang, the high-performance serving infrastructure powering trillions of tokens daily, has been awarded the third batch of the Open Source AI Grant by a16z ([a16z blog](https://a16z.com/advancing-open-source-ai-through-benchmarks-and-bold-experimentation/)). +- [2025/06] 🔥 Deploying DeepSeek on GB200 NVL72 with PD and Large Scale EP (Part I): 2.7x Higher Decoding Throughput ([blog](https://lmsys.org/blog/2025-06-16-gb200-part-1/)). +- [2025/05] 🔥 Deploying DeepSeek with PD Disaggregation and Large-scale Expert Parallelism on 96 H100 GPUs ([blog](https://lmsys.org/blog/2025-05-05-large-scale-ep/)). +- [2025/03] Supercharge DeepSeek-R1 Inference on AMD Instinct MI300X ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1-Part2/README.html)) +- [2025/03] SGLang Joins PyTorch Ecosystem: Efficient LLM Serving Engine ([PyTorch blog](https://pytorch.org/blog/sglang-joins-pytorch/)) +- [2024/12] v0.4 Release: Zero-Overhead Batch Scheduler, Cache-Aware Load Balancer, Faster Structured Outputs ([blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/)). + +
+More + +- [2025/02] Unlock DeepSeek-R1 Inference Performance on AMD Instinct™ MI300X GPU ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1_Perf/README.html)) +- [2025/01] SGLang provides day one support for DeepSeek V3/R1 models on NVIDIA and AMD GPUs with DeepSeek-specific optimizations. ([instructions](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3), [AMD blog](https://www.amd.com/en/developer/resources/technical-articles/amd-instinct-gpus-power-deepseek-v3-revolutionizing-ai-development-with-sglang.html), [10+ other companies](https://x.com/lmsysorg/status/1887262321636221412)) +- [2024/10] The First SGLang Online Meetup ([slides](https://github.com/sgl-project/sgl-learning-materials?tab=readme-ov-file#the-first-sglang-online-meetup)). +- [2024/09] v0.3 Release: 7x Faster DeepSeek MLA, 1.5x Faster torch.compile, Multi-Image/Video LLaVA-OneVision ([blog](https://lmsys.org/blog/2024-09-04-sglang-v0-3/)). +- [2024/07] v0.2 Release: Faster Llama3 Serving with SGLang Runtime (vs. TensorRT-LLM, vLLM) ([blog](https://lmsys.org/blog/2024-07-25-sglang-llama3/)). +- [2024/02] SGLang enables **3x faster JSON decoding** with compressed finite state machine ([blog](https://lmsys.org/blog/2024-02-05-compressed-fsm/)). +- [2024/01] SGLang provides up to **5x faster inference** with RadixAttention ([blog](https://lmsys.org/blog/2024-01-17-sglang/)). +- [2024/01] SGLang powers the serving of the official **LLaVA v1.6** release demo ([usage](https://github.com/haotian-liu/LLaVA?tab=readme-ov-file#demo)). + +
+ +## About +SGLang is a fast serving framework for large language models and vision language models. +It makes your interaction with models faster and more controllable by co-designing the backend runtime and frontend language. +The core features include: + +- **Fast Backend Runtime**: Provides efficient serving with RadixAttention for prefix caching, zero-overhead CPU scheduler, prefill-decode disaggregation, speculative decoding, continuous batching, paged attention, tensor/pipeline/expert/data parallelism, structured outputs, chunked prefill, quantization (FP4/FP8/INT4/AWQ/GPTQ), and multi-lora batching. +- **Flexible Frontend Language**: Offers an intuitive interface for programming LLM applications, including chained generation calls, advanced prompting, control flow, multi-modal inputs, parallelism, and external interactions. +- **Extensive Model Support**: Supports a wide range of generative models (Llama, Qwen, DeepSeek, Kimi, GPT, Gemma, Mistral, etc.), embedding models (e5-mistral, gte, mcdse) and reward models (Skywork), with easy extensibility for integrating new models. +- **Active Community**: SGLang is open-source and backed by an active community with wide industry adoption. + +## Getting Started +- [Install SGLang](https://docs.sglang.ai/get_started/install.html) +- [Quick Start](https://docs.sglang.ai/basic_usage/send_request.html) +- [Backend Tutorial](https://docs.sglang.ai/basic_usage/openai_api_completions.html) +- [Frontend Tutorial](https://docs.sglang.ai/references/frontend/frontend_tutorial.html) +- [Contribution Guide](https://docs.sglang.ai/developer_guide/contribution_guide.html) + +## Benchmark and Performance +Learn more in the release blogs: [v0.2 blog](https://lmsys.org/blog/2024-07-25-sglang-llama3/), [v0.3 blog](https://lmsys.org/blog/2024-09-04-sglang-v0-3/), [v0.4 blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/), [Large-scale expert parallelism](https://lmsys.org/blog/2025-05-05-large-scale-ep/). + +## Roadmap +[Development Roadmap (2025 H2)](https://github.com/sgl-project/sglang/issues/7736) + +## Adoption and Sponsorship +SGLang has been deployed at large scale, generating trillions of tokens in production each day. It is trusted and adopted by a wide range of leading enterprises and institutions, including xAI, AMD, NVIDIA, Intel, LinkedIn, Cursor, Oracle Cloud, Google Cloud, Microsoft Azure, AWS, Atlas Cloud, Voltage Park, Nebius, DataCrunch, Novita, InnoMatrix, MIT, UCLA, the University of Washington, Stanford, UC Berkeley, Tsinghua University, Jam & Tea Studios, Baseten, and other major technology organizations across North America and Asia. As an open-source LLM inference engine, SGLang has become the de facto industry standard, with deployments running on over 1,000,000 GPUs worldwide. + +logo + +## Contact Us +For enterprises interested in adopting or deploying SGLang at scale, including technical consulting, sponsorship opportunities, or partnership inquiries, please contact us at contact@sglang.ai. + +## Acknowledgment +We learned the design and reused code from the following projects: [Guidance](https://github.com/guidance-ai/guidance), [vLLM](https://github.com/vllm-project/vllm), [LightLLM](https://github.com/ModelTC/lightllm), [FlashInfer](https://github.com/flashinfer-ai/flashinfer), [Outlines](https://github.com/outlines-dev/outlines), and [LMQL](https://github.com/eth-sri/lmql). diff --git a/assets/logo.png b/assets/logo.png new file mode 100644 index 000000000..2a8bc258f Binary files /dev/null and b/assets/logo.png differ diff --git a/assets/logo.svg b/assets/logo.svg new file mode 100644 index 000000000..4d6393926 --- /dev/null +++ b/assets/logo.svg @@ -0,0 +1 @@ + diff --git a/assets/logo_square.svg b/assets/logo_square.svg new file mode 100644 index 000000000..a82fa0aeb --- /dev/null +++ b/assets/logo_square.svg @@ -0,0 +1 @@ + diff --git a/benchmark/bench_attention_sink/bench_attention_sink_triton.py b/benchmark/bench_attention_sink/bench_attention_sink_triton.py new file mode 100644 index 000000000..21bc2a593 --- /dev/null +++ b/benchmark/bench_attention_sink/bench_attention_sink_triton.py @@ -0,0 +1,250 @@ +import argparse + +import torch +import triton + +from sglang.srt.layers.attention.triton_ops.decode_attention import ( + decode_attention_fwd_grouped, +) +from sglang.srt.layers.attention.triton_ops.extend_attention import extend_attention_fwd + +# gpt oss +head_num = 64 +head_dim = 64 +head_kv_num = 8 + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["S"], # sequence length on x-axis + x_vals=[128, 256, 512, 1024, 2048, 4096], + x_log=True, + line_arg="B", # batch size as different lines + line_vals=[1, 8, 32, 128], + line_names=["B=1", "B=8", "B=32", "B=128"], + styles=[ + ("blue", "-"), + ("green", "-"), + ("red", "-"), + ("cyan", "-"), + ], + ylabel="TFLOPS", + plot_name="attention-sink-triton-decode", + args={}, + ) +) +def benchmark_decode(B, S, H_Q, H_KV, D): + D_V = D + dtype = torch.bfloat16 + seq_len = S + total_tokens = B * seq_len + device = torch.device("cuda") + sm_scale = 1.0 / (D**0.5) + max_kv_splits = 8 + num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device="cuda") + + # q represents the new token being generated, one per batch + q = torch.randn(B, H_Q, D, dtype=dtype, device="cuda") + + # k_buffer and v_buffer represent all previous tokens + k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") + v_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") + + o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda") + + b_seq_len = torch.full((B,), seq_len, device="cuda") + + kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") + kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len, dim=0) + kv_indices = torch.arange(total_tokens, device="cuda") + + attn_logits1 = torch.empty( + (B, H_Q, max_kv_splits, D_V), + dtype=torch.float32, + device="cuda", + ) + attn_lse1 = torch.empty( + (B, H_Q, max_kv_splits, D_V), + dtype=torch.float32, + device="cuda", + ) + sink = torch.randn(H_Q, device=device, dtype=torch.float32) + + # 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 + run_step = 500 + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + for _ in range(run_step): + 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, + ) + end_event.record() + end_event.synchronize() + torch.cuda.synchronize() + ms = start_event.elapsed_time(end_event) / run_step + tflops = lambda ms: (2 * B * S * H_Q * D) * 1e-9 / ms # must be causal + return tflops(ms) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["S"], # sequence length on x-axis + x_vals=[128, 256, 512, 1024, 2048, 4096], + x_log=True, + line_arg="B", # batch size as different lines + line_vals=[1, 8, 32, 128], + line_names=["B=1", "B=8", "B=32", "B=128"], + styles=[ + ("blue", "-"), + ("green", "-"), + ("red", "-"), + ("cyan", "-"), + ], + ylabel="TFLOPS", + plot_name="attention-sink-triton-extend", + args={}, + ) +) +def benchmark_extend(B, S, H_Q, H_KV, D): + # S here represents N_CTX from the test + dtype = torch.bfloat16 + device = "cuda" + + # Split S into prefix and extend lengths + prefill_len = S // 2 # Similar to test's N_CTX // 2 + extend_len = S // 4 # Make extend length smaller than prefix + + # Calculate total tokens and extend tokens + total_extend_tokens = B * extend_len + total_prefix_tokens = B * prefill_len + + # Create query, key, value tensors for extension + q_extend = torch.randn(total_extend_tokens, H_Q, D, dtype=dtype, device=device) + k_extend = torch.randn(total_extend_tokens, H_KV, D, dtype=dtype, device=device) + v_extend = torch.randn(total_extend_tokens, H_KV, D, dtype=dtype, device=device) + o_extend = torch.empty_like(q_extend) + + # Create key-value buffers for prefix + k_buffer = torch.randn(total_prefix_tokens, H_KV, D, dtype=dtype, device=device) + v_buffer = torch.randn(total_prefix_tokens, H_KV, D, dtype=dtype, device=device) + + # Create index pointers + qo_indptr = torch.arange(0, (B + 1) * extend_len, extend_len, device=device).to( + torch.int32 + ) + kv_indptr = torch.arange(0, (B + 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 / (D**0.5) + # sliding_window = 128 # From GPT-OSS config, skip for now + sliding_window = -1 + + sink = torch.randn(H_Q, device=device, dtype=torch.float32) + + # 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 + run_step = 500 + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + for _ in range(run_step): + 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, + ) + end_event.record() + end_event.synchronize() + torch.cuda.synchronize() + ms = start_event.elapsed_time(end_event) / run_step + + # FLOPS calculation: each attention operation requires 2 multiplications per element + total_flops = 2 * total_extend_tokens * H_Q * (prefill_len + extend_len / 2) * D + tflops = lambda ms: total_flops * 1e-12 / (ms * 1e-3) # convert to TFLOPS + return tflops(ms) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--bench", type=str, default="all", help="all, extend, decode") + args = parser.parse_args() + + kwargs = { + "H_Q": head_num, + "H_KV": head_kv_num, + "D": head_dim, + } + + if args.bench in ["all", "decode"]: + benchmark_decode.run(print_data=True, show_plots=False, **kwargs) + + if args.bench in ["all", "extend"]: + benchmark_extend.run(print_data=True, show_plots=False, **kwargs) + + print("Benchmark finished!") diff --git a/benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py b/benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py new file mode 100644 index 000000000..282097112 --- /dev/null +++ b/benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py @@ -0,0 +1,130 @@ +# Benchmark with lots of common prefixes. Used to benchmark prefix caching performance. +# +# Launch a server: +# python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --log-level-http warning + +import random +import string +import time + +from tqdm import tqdm +from transformers import AutoTokenizer + +import sglang as sgl +from sglang import set_default_backend +from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint + + +def generate_random_string(token_length: int) -> str: + random_string = "".join( + random.choices(string.ascii_letters + string.digits, k=token_length * 100) + ) + tokenized_output = tokenizer.encode(random_string, add_special_tokens=False)[ + :token_length + ] + + if len(tokenized_output) < token_length: + tokenized_output = tokenized_output + [tokenizer.pad_token_id] * ( + token_length - len(tokenized_output) + ) + + decoded_string = tokenizer.decode(tokenized_output, skip_special_tokens=False) + return decoded_string + + +def generate_unique_prefix(base_text, index): + return str(index) + base_text[len(str(index)) :] + + +@sgl.function +def text_qa(s, question, gen_len): + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0, max_tokens=gen_len) + + +def prepare_prompts(num_prefix, num_samples_per_prefix, prefix_length, suffix_length): + base_prefix = generate_random_string(prefix_length) + + tot_input_len = 0 + all_prompts = [] + for i in tqdm(range(num_prefix), desc="prepare prompts"): + unique_prefix = generate_unique_prefix(base_prefix, i) + prompt_list = [] + for j in range(num_samples_per_prefix): + suffix = generate_random_string(suffix_length) + prompt = unique_prefix + suffix + prompt_list.append(prompt) + tot_input_len += len(tokenizer.encode(prompt)) + all_prompts.append(prompt_list) + return all_prompts, tot_input_len + + +def test_batch_by_batch(all_prompts, gen_len): + backend.flush_cache() + + tot_time = 0 + for i in range(len(all_prompts)): + tic = time.perf_counter() + text_qa.run_batch( + list(zip(all_prompts[i], [gen_len] * len(all_prompts[i]))), + ) + tot_time += time.perf_counter() - tic + + return tot_time + + +def test_batch_by_batch_with_hint(all_prompts, gen_len): + backend.flush_cache() + + tot_time = 0 + for i in range(len(all_prompts)): + tic = time.perf_counter() + # Send a hint to cache the prefix + text_qa.run_batch(list(zip(all_prompts[i][:1], [gen_len]))) + # Send the batch + text_qa.run_batch(list(zip(all_prompts[i], [gen_len] * len(all_prompts[i])))) + + tot_time += time.perf_counter() - tic + + return tot_time + + +def test_send_all(all_prompts, gen_len): + backend.flush_cache() + + all_prompts = [x for prompt_list in all_prompts for x in prompt_list] + + tic = time.perf_counter() + text_qa.run_batch( + list(zip(all_prompts, [gen_len] * len(all_prompts))), + ) + tot_time = time.perf_counter() - tic + + return tot_time + + +if __name__ == "__main__": + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") + backend = RuntimeEndpoint("http://127.0.0.1:30000") + set_default_backend(backend) + + random.seed(0) + num_prefix = 10 + num_samples_per_prefix = 32 + prefix_length = 1024 + suffix_length = 128 + gen_len = 1 + all_prompts, tot_input_len = prepare_prompts( + num_prefix, num_samples_per_prefix, prefix_length, suffix_length + ) + + print(f"Total input token length: {tot_input_len}\n") + + cost = test_batch_by_batch(all_prompts, gen_len) + print(f"Latency of test_batch_by_batch : {cost:.4f} s\n") + + cost = test_batch_by_batch_with_hint(all_prompts, gen_len) + print(f"Latency of test_batch_by_batch_with_hint: {cost:.4f} s\n") + + cost = test_send_all(all_prompts, gen_len) + print(f"Latency of test_send_all : {cost:.4f} s\n") diff --git a/benchmark/benchmark_batch/benchmark_batch.py b/benchmark/benchmark_batch/benchmark_batch.py new file mode 100644 index 000000000..a8592d48a --- /dev/null +++ b/benchmark/benchmark_batch/benchmark_batch.py @@ -0,0 +1,193 @@ +import concurrent.futures +import os +import random +import time +from concurrent.futures import ProcessPoolExecutor +from statistics import mean + +import requests +from tqdm import tqdm +from transformers import AutoTokenizer + +from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint + +############################################################################### +# CONFIG +############################################################################### +ENDPOINT_URL = "http://127.0.0.1:30000" +TOKENIZER_DIR = "/models/meta-llama/Llama-3.2-3B" + +# Benchmark configurations +NUM_REQUESTS = 10 # Total number of requests (each with BATCH_SIZE prompts) +NUM_TOKENS = 32000 # Tokens per prompt +BATCH_SIZE = 8 # Number of prompts per request +GEN_TOKENS = 0 # Tokens to generate per prompt + + +############################################################################### +# REQUEST GENERATION (in parallel) +############################################################################### +def generate_random_prompt(index, tokenizer_dir, num_tokens): + """Generate a single random prompt with specified token count.""" + tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir) + vocab_size = tokenizer.vocab_size + + def generate_random_text(num_toks): + random_token_ids = [random.randint(0, vocab_size - 1) for _ in range(num_toks)] + return tokenizer.decode(random_token_ids, clean_up_tokenization_spaces=True) + + random_text = generate_random_text(num_tokens) + return f"Prompt {index}: {random_text}" + + +def prepare_all_prompts(num_requests, batch_size, num_tokens, tokenizer_dir): + """Generate prompts for all requests in parallel.""" + total_prompts = num_requests * batch_size + all_prompts = [None] * total_prompts + max_workers = min(os.cpu_count() or 1, total_prompts) + + with ProcessPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit(generate_random_prompt, i, tokenizer_dir, num_tokens) + for i in range(total_prompts) + ] + for future in tqdm( + concurrent.futures.as_completed(futures), + total=total_prompts, + desc="Generating prompts", + ): + index = futures.index(future) + all_prompts[index] = future.result() + + batched_prompts = [ + all_prompts[i * batch_size : (i + 1) * batch_size] for i in range(num_requests) + ] + + print( + f"Generated {total_prompts} prompts with {num_tokens} tokens each, grouped into {num_requests} requests of {batch_size} prompts.\n" + ) + return batched_prompts + + +############################################################################### +# HTTP CALLS +############################################################################### +def send_batch_request(endpoint, prompts, gen_tokens, request_id): + """Send a batch of prompts to the /generate endpoint synchronously.""" + sampling_params = { + "max_new_tokens": gen_tokens, + "temperature": 0.7, + "stop": "\n", + } + data = {"text": prompts, "sampling_params": sampling_params} + + start_time = time.perf_counter() + try: + response = requests.post( + endpoint.base_url + "/generate", json=data, timeout=3600 + ) + if response.status_code != 200: + error = response.json() + raise RuntimeError(f"Request {request_id} failed: {error}") + result = response.json() + elapsed_time = (time.perf_counter() - start_time) * 1000 # Convert to ms + avg_per_prompt = elapsed_time / len(prompts) if prompts else 0 + return request_id, elapsed_time, avg_per_prompt, True, len(prompts) + except Exception as e: + print(f"[Request] Error for request {request_id}: {e}") + return request_id, 0, 0, False, len(prompts) + + +def run_benchmark(endpoint, batched_prompts, batch_size, gen_tokens): + """Run the benchmark sequentially.""" + results = [] + num_requests = len(batched_prompts) + + # Record start time for total latency + benchmark_start_time = time.perf_counter() + + for i, batch_prompts in enumerate(batched_prompts): + request_id = i + 1 + assert ( + len(batch_prompts) == batch_size + ), f"Request {request_id} should have {batch_size} prompts, got {len(batch_prompts)}" + + print( + f"[Request] Sending request {request_id}/{num_requests} with {len(batch_prompts)} prompts at {int(time.time()*1000)}" + ) + result = send_batch_request(endpoint, batch_prompts, gen_tokens, request_id) + results.append(result) + + # Calculate total latency + total_latency = (time.perf_counter() - benchmark_start_time) * 1000 # Convert to ms + + return results, total_latency + + +############################################################################### +# RESULTS +############################################################################### +def process_results(results, total_latency, num_requests): + """Process and display benchmark results.""" + total_time = 0 + successful_requests = 0 + failed_requests = 0 + request_latencies = [] + per_prompt_latencies = [] + total_prompts = 0 + + for request_id, elapsed_time, avg_per_prompt, success, batch_size in results: + if success: + successful_requests += 1 + total_prompts += batch_size + request_latencies.append(elapsed_time) + per_prompt_latencies.append(avg_per_prompt) + total_time += elapsed_time / 1000 # Convert to seconds + else: + failed_requests += 1 + + avg_request_latency = mean(request_latencies) if request_latencies else 0 + avg_per_prompt_latency = mean(per_prompt_latencies) if per_prompt_latencies else 0 + throughput = total_prompts / total_time if total_time > 0 else 0 + + print("\nBenchmark Summary:") + print(f" Total requests sent: {len(results)}") + print(f" Total prompts sent: {total_prompts}") + print(f" Successful requests: {successful_requests}") + print(f" Failed requests: {failed_requests}") + print(f" Total latency (all requests): {total_latency:.2f} ms") + print(f" Avg per request latency: {avg_request_latency:.2f} ms") + print(f" Avg per prompt latency: {avg_per_prompt_latency:.2f} ms") + print(f" Throughput: {throughput:.2f} prompts/second\n") + + +############################################################################### +# MAIN +############################################################################### +def main(): + # Initialize endpoint + endpoint = RuntimeEndpoint(ENDPOINT_URL) + + # Generate prompts + batched_prompts = prepare_all_prompts( + NUM_REQUESTS, BATCH_SIZE, NUM_TOKENS, TOKENIZER_DIR + ) + + # Flush cache before benchmark + # endpoint.flush_cache() + + # Run benchmark + print( + f"Starting benchmark: NUM_TOKENS={NUM_TOKENS}, BATCH_SIZE={BATCH_SIZE}, NUM_REQUESTS={NUM_REQUESTS}\n" + ) + results, total_latency = run_benchmark( + endpoint, batched_prompts, BATCH_SIZE, GEN_TOKENS + ) + + # Process and display results + process_results(results, total_latency, NUM_REQUESTS) + + +if __name__ == "__main__": + random.seed(0) + main() diff --git a/benchmark/benchmark_batch/benchmark_tokenizer.py b/benchmark/benchmark_batch/benchmark_tokenizer.py new file mode 100644 index 000000000..88a5820b6 --- /dev/null +++ b/benchmark/benchmark_batch/benchmark_tokenizer.py @@ -0,0 +1,126 @@ +import random +import time +from statistics import mean + +from transformers import AutoTokenizer + +# CONFIG +TOKENIZER_DIR = ( + "/shared/public/sharing/fait360brew/training/models/meta-llama/Llama-3.2-3B" +) +NUM_TOKENS = 20000 # Each prompt should contain this many tokens +BATCH_SIZES = [1, 2, 4, 8] # Test different batch sizes +NUM_RUNS = 5 # Number of runs for each batch size to get reliable measurements + + +def generate_random_prompts(num_prompts, num_tokens, tokenizer): + """Generate random prompts with specified token count.""" + vocab_size = tokenizer.vocab_size + all_prompts = [] + + print(f"Generating {num_prompts} random prompts with {num_tokens} tokens each...") + for i in range(num_prompts): + # Generate random token IDs - this directly gives us the exact token count + random_token_ids = [ + random.randint(0, vocab_size - 1) for _ in range(num_tokens) + ] + random_text = tokenizer.decode( + random_token_ids, clean_up_tokenization_spaces=True + ) + + prompt = f"Prompt {i}: {random_text}" + tokens = tokenizer.encode(prompt) + print(f" Prompt {i}: {len(tokens)} tokens") + all_prompts.append(prompt) + + return all_prompts + + +def benchmark_sequential_vs_batch(prompts, batch_size, tokenizer): + """Compare sequential vs batch tokenization for a given batch size.""" + + # Sequential tokenization using encode() + sequential_times = [] + for run in range(NUM_RUNS): + batch_prompts = prompts[:batch_size] # Use same prompts for fair comparison + + start_time = time.perf_counter() + for prompt in batch_prompts: + tokens = tokenizer.encode(prompt) + sequential_time = (time.perf_counter() - start_time) * 1000 + sequential_times.append(sequential_time) + + # Batch tokenization using tokenizer() + batch_times = [] + for run in range(NUM_RUNS): + batch_prompts = prompts[:batch_size] # Use same prompts for fair comparison + + start_time = time.perf_counter() + tokens = tokenizer(batch_prompts) + batch_time = (time.perf_counter() - start_time) * 1000 + batch_times.append(batch_time) + + return { + "batch_size": batch_size, + "avg_sequential_ms": mean(sequential_times), + "avg_batch_ms": mean(batch_times), + "speedup_factor": ( + mean(sequential_times) / mean(batch_times) if mean(batch_times) > 0 else 0 + ), + "sequential_runs": sequential_times, + "batch_runs": batch_times, + } + + +def main(): + print("Tokenizer Benchmark: Sequential vs Batch Processing") + print("-" * 60) + print(f"Tokenizer: {TOKENIZER_DIR}") + print(f"Tokens per prompt: {NUM_TOKENS}") + print(f"Number of runs per batch size: {NUM_RUNS}") + print("-" * 60) + + # Load tokenizer once for all operations + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR) + + # The largest batch size determines how many prompts we need + max_batch_size = max(BATCH_SIZES) + all_prompts = generate_random_prompts(max_batch_size, NUM_TOKENS, tokenizer) + + results = [] + print("\nRunning benchmark...") + + for batch_size in BATCH_SIZES: + print(f"\nBenchmarking batch size: {batch_size}") + result = benchmark_sequential_vs_batch(all_prompts, batch_size, tokenizer) + results.append(result) + + print(f" Sequential tokenization (encode):") + for i, run_time in enumerate(result["sequential_runs"]): + print(f" Run {i+1}: {run_time:.2f} ms") + print(f" Average: {result['avg_sequential_ms']:.2f} ms") + + print(f" Batch tokenization (tokenizer):") + for i, run_time in enumerate(result["batch_runs"]): + print(f" Run {i+1}: {run_time:.2f} ms") + print(f" Average: {result['avg_batch_ms']:.2f} ms") + + print(f" Speedup factor: {result['speedup_factor']:.2f}x") + + print("\n" + "=" * 60) + print("SUMMARY OF RESULTS") + print("=" * 60) + print( + f"{'Batch Size':<10} {'Sequential (ms)':<18} {'Batch (ms)':<18} {'Speedup':<10}" + ) + print("-" * 60) + + for result in results: + print( + f"{result['batch_size']:<10} {result['avg_sequential_ms']:.2f} ms{' ' * 8} {result['avg_batch_ms']:.2f} ms{' ' * 8} {result['speedup_factor']:.2f}x" + ) + + +if __name__ == "__main__": + random.seed(0) + main() diff --git a/benchmark/benchmark_vllm_060/README.md b/benchmark/benchmark_vllm_060/README.md new file mode 100644 index 000000000..b480dabf2 --- /dev/null +++ b/benchmark/benchmark_vllm_060/README.md @@ -0,0 +1,89 @@ +## How to reproduce the benchmark results for SGLang v0.3.0 compared to vLLM v0.6.0 + +In short, with multi step enabled, in online scenarios that we benchmarked, the Median TTFT of vLLM is **3 times** that of SGLang, and the Median ITL is **10 times** that of SGLang. Lower Median TTFT and ITL are better. vLLM's multi-step optimization did not improve throughput while ensuring lower Median TTFT and ITL. Also, under maximum throughput benchmark, if vLLM does not set gpu util to 0.95 separately and uses the default configuration instead, its maximum throughput is **lower** than that of SGLang. + +## Online benchmark results + +### Llama 3.1 8B Instruct 1 x A100 80G + +| RPS | Num prompts | Engine | Median E2E Latency | Median TTFT | Median TPOT | Median ITL | +|------|-------------|--------|--------------------|-------------|-------------|------------| +| 4 | 1200 | SGLang | 1564.17 | **31.98** | 13.17 | **11.93** | +| 4 | 1200 | vLLM | 1691.97 | **100.48** | 14.14 | **129.32** | +| 8 | 2400 | SGLang | 2175.02 | **35.68** | 17.85 | **14.41** | +| 8 | 2400 | vLLM | 2137.16 | **120.39** | 17.09 | **158.63** | + +### Llama 3.1 70B Insruct 4 x H100 80G + +| RPS | Num Prompts | Engine | Median E2E Latency | Median TTFT | Median TPOT | Median ITL | +|------|-------------|--------|--------------------|-------------|-------------|------------| +| 4 | 1200 | SGLang | 3005.24 | **53.94** | 25.03 | **21.67** | +| 4 | 1200 | vLLM | 2915.60 | **179.15** | 23.58 | **231.23** | +| 8 | 2400 | SGLang | 4064.98 | **58.11** | 33.07 | **24.45** | +| 8 | 2400 | vLLM | 3752.38 | **207.12** | 29.15 | **275.32** | + +## Offline benchmark results + +### Llama 3.1 8B Instruct 1 x A100 80G + +| RPS | Num Prompts | Engine | Request throughput | Output token throughput | +|------|-------------|--------|--------------------|-------------------------| +| inf | 5000 | SGLang | 22.03 | **4281.51** | +| inf | 5000 | vLLM | 21.27 | **4132.37** | + +### Llama 3.1 70B Insruct 4 x H100 80G + +| RPS | Num Prompts | Engine | Request throughput | Output token throughput | +|------|-------------|--------|--------------------|-------------------------| +| inf | 5000 | SGLang | 19.84 | **3856.01** | +| inf | 5000 | vLLM | 19.04 | **3700.64** | + +## Installation + +```bash +# install sglang v0.3.0 +pip install --upgrade pip +pip install "sglang[all]"==0.3.0 +pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.4/ + +# install vllm v0.6.0 +pip install vllm==0.6.0 +``` + +## Notes + +We referred to the reproduction method in https://github.com/vllm-project/vllm/issues/8176, and added the `--num-scheduler-steps 10` parameter when starting the vLLM server. The `gpu_memory_utilization` of vLLM is by default 0.9 at both TP 1 and TP 4, while SGLang's `mem_frac` is 0.88 at TP 1 and 0.85 at TP 4, so we manually set it to 0.88 at TP 4. + +## Online benchmarks + +```bash +# Llama 3.1 8B Instruct on 1 x A100 +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile --disable-radix-cache +python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-8B-Instruct --disable-log-requests --num-scheduler-steps 10 --max_model_len 4096 + +# Llama 3.1 70B Instruct on 4 x H100 +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-70B-Instruct --disable-radix-cache --tp 4 +python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-70B-Instruct --disable-log-requests --num-scheduler-steps 10 --tensor 4 --max_model_len 4096 + +# bench serving +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 1200 --request-rate 4 +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 2400 --request-rate 8 +python3 -m sglang.bench_serving --backend vllm --dataset-name sharegpt --num-prompts 1200 --request-rate 4 +python3 -m sglang.bench_serving --backend vllm --dataset-name sharegpt --num-prompts 2400 --request-rate 8 +``` + +## Offline benchmarks + +```bash +# Llama 3.1 8B Instruct on 1 x A100 +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile --disable-radix-cache +python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-8B-Instruct --disable-log-requests --num-scheduler-steps 10 --max_model_len 4096 + +# Llama 3.1 70B Instruct on 4 x H100 +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-70B-Instruct --disable-radix-cache --tp 4 --mem-frac 0.88 +python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-70B-Instruct --disable-log-requests --num-scheduler-steps 10 --tensor 4 --max_model_len 4096 + +# bench serving +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 5000 +python3 -m sglang.bench_serving --backend vllm --dataset-name sharegpt --num-prompts 5000 +``` diff --git a/benchmark/blog_v0_2/405b_sglang.sh b/benchmark/blog_v0_2/405b_sglang.sh new file mode 100644 index 000000000..491853782 --- /dev/null +++ b/benchmark/blog_v0_2/405b_sglang.sh @@ -0,0 +1,24 @@ +# Create dummy weights: +# 1. Create a folder `~/llama-3.1-405b-fp8-dummy` and create `config.json` and tokenizer under this folder. +# 2. Get `config.json`` from ./config.md +# 3. Download the tokenizer +# wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer.json +# wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer_config.json + +# Launch sglang +# python -m sglang.launch_server --model-path ~/llama-3.1-405b-fp8-dummy/ --load-format dummy --tp 8 --quantization fp8 --disable-radix --mem-frac 0.87 + +# offline +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 3000 --random-input 1024 --random-output 1024 > sglang_log11 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 4000 --random-input 1024 --random-output 512 > sglang_log12 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 800 --random-input 4096 --random-output 2048 > sglang_log13 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 1500 --random-input 4096 --random-output 1024 > sglang_log14 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 6000 --random-input 256 --random-output 512 > sglang_log15 +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompt 2000 > sglang_log21 + +# online +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 300 --request-rate 1 --random-input 1024 --random-output 1024 > sglang_log31 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 600 --request-rate 2 --random-input 1024 --random-output 1024 > sglang_log32 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 1200 --request-rate 4 --random-input 1024 --random-output 1024 > sglang_log33 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 2400 --request-rate 8 --random-input 1024 --random-output 1024 > sglang_log34 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 3200 --request-rate 16 --random-input 1024 --random-output 1024 > sglang_log35 diff --git a/benchmark/blog_v0_2/405b_trt.sh b/benchmark/blog_v0_2/405b_trt.sh new file mode 100644 index 000000000..1950bc92b --- /dev/null +++ b/benchmark/blog_v0_2/405b_trt.sh @@ -0,0 +1,17 @@ +# Launch trtllm +# https://github.com/sgl-project/tensorrt-demo + +# offline +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 3000 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log11 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 4000 --random-input 1024 --random-output 512 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log12 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 800 --random-input 4096 --random-output 2048 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log13 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 1500 --random-input 4096 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log14 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 6000 --random-input 256 --random-output 512 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log15 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name sharegpt --num-prompt 2000 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log21 + +# online +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 300 --request-rate 1 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log31 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 600 --request-rate 2 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log32 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 1200 --request-rate 4 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log33 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 2400 --request-rate 8 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log34 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 3200 --request-rate 16 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log35 diff --git a/benchmark/blog_v0_2/405b_vllm.sh b/benchmark/blog_v0_2/405b_vllm.sh new file mode 100644 index 000000000..d1ef048e8 --- /dev/null +++ b/benchmark/blog_v0_2/405b_vllm.sh @@ -0,0 +1,24 @@ +# Create dummy weights: +# 1. Create a folder `~/llama-3.1-405b-fp8-dummy` and create `config.json` and tokenizer under this folder. +# 2. Get `config.json`` from ./config.md +# 3. Download the tokenizer +# wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer.json +# wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer_config.json + +# Launch vllm +# python3 -m vllm.entrypoints.openai.api_server --model ~/llama-3.1-405b-fp8-dummy/ --load-format dummy --disable-log-requests --tensor-parallel-size 8 --max-model-len 10000 + +# offline +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 3000 --random-input 1024 --random-output 1024 > vllm_log11 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 4000 --random-input 1024 --random-output 512 > vllm_log12 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 800 --random-input 4096 --random-output 2048 > vllm_log13 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 1500 --random-input 4096 --random-output 1024 > vllm_log14 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 6000 --random-input 256 --random-output 512 > vllm_log15 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name sharegpt --num-prompt 2000 > vllm_log21 + +# online +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 300 --request-rate 1 --random-input 1024 --random-output 1024 > vllm_log31 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 600 --request-rate 2 --random-input 1024 --random-output 1024 > vllm_log32 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 1200 --request-rate 4 --random-input 1024 --random-output 1024 > vllm_log33 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 2400 --request-rate 8 --random-input 1024 --random-output 1024 > vllm_log34 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 3200 --request-rate 16 --random-input 1024 --random-output 1024 > vllm_log35 diff --git a/benchmark/blog_v0_2/README.md b/benchmark/blog_v0_2/README.md new file mode 100644 index 000000000..7448554ee --- /dev/null +++ b/benchmark/blog_v0_2/README.md @@ -0,0 +1,164 @@ +# How to reproduce the benchmark results of SGLang + +## Prerequisite + +### Install the latest SGLang + +```bash +git clone https://github.com/sgl-project/sglang.git +cd sglang +git checkout v0.2.7 + +pip install --upgrade pip +pip install -e "python[all]" + +pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.3/ +``` + +### Set up ulimit and HF_TOKEN + +```bash +ulimit -n 65535 +# Change the token to a real and usable one, with access permissions for the Llama 3 models. +export HF_TOKEN=hf_token +``` + +### Launch the server + +```bash +# Meta-Llama-3.1-8B-Instruct +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile --disable-radix-cache + +# Meta-Llama-3.1-70B-Instruct +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-70B-Instruct --disable-radix-cache --tp 8 + +# Meta-Llama-3-70B-Instruct-FP8 +python -m sglang.launch_server --model-path neuralmagic/Meta-Llama-3-70B-Instruct-FP8 --disable-radix-cache --tp 8 +``` + +## Benchmark + +### Hardware Requirements + +- 8B models: Single NVIDIA A100 80GB GPU +- 70B models: 8 x NVIDIA A100 80GB GPUs with Tensor Parallelism (TP) 8 +- 70B FP8 models: 8 x NVIDIA H100 GPUs with Tensor Parallelism (TP) 8 + +Please ensure you have the appropriate hardware before running the benchmarks. + +#### Offline benchmark + +```bash +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 3000 --output-file offline.jsonl +cat offline.jsonl | cut -d':' -f12 | cut -d',' -f1 +``` + +#### Online benchmark + +```bash +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online.jsonl +cat online.jsonl | cut -d':' -f9 | cut -d',' -f1 +``` + +## Other + +We tried using vLLM 0.5.3.post1, but it often crashes under high loads, and it seems to have similar or worse performance compared to vLLM 0.5.2 from our partial benchmarking, so we are using the older version, vLLM 0.5.2. + +Preparation for TensorRT LLM can refer to https://github.com/sgl-project/tensorrt-demo. Specifically, we used a batch size of 512, a max input length of 8192, and a max number of tokens of 8192. The instance count for preprocessing and postprocessing in Triton Server is 16. + +```bash +# vLLM +pip install vllm==0.5.2 +pip install jsonschema==4.21.1 + +# Meta-Llama-3-8B-Instruct +python -m vllm.entrypoints.openai.api_server --model meta-llama/Meta-Llama-3-8B-Instruct --disable-log-requests + +# meta-llama/Meta-Llama-3-70B-Instruct +python -m vllm.entrypoints.openai.api_server --model meta-llama/Meta-Llama-3-70B-Instruct --disable-log-requests --tensor 8 + +# neuralmagic/Meta-Llama-3-70B-Instruct-FP8 +python -m vllm.entrypoints.openai.api_server --model neuralmagic/Meta-Llama-3-70B-Instruct-FP8 --disable-log-requests --tensor 8 +``` + +```bash +wget https://raw.githubusercontent.com/sgl-project/sglang/main/python/sglang/bench_serving.py +``` + +```bash +# vLLM Offline + +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name sharegpt --num-prompts 3000 --output-file offline_vllm.jsonl +cat offline_vllm.jsonl | cut -d':' -f12 | cut -d',' -f1 +``` + +```bash +# vLLM Online + +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online_vllm.jsonl +cat online_vllm.jsonl | cut -d':' -f9 | cut -d',' -f1 +``` + +```bash +# TensorRT LLM Offline 8B + +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline_trt_8b.jsonl +python3 bench_serving.py --backend trt --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline_trt_8b.jsonl --model meta-llama/Meta-Llama-3-8B-Instruct +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name sharegpt --num-prompts 3000 --output-file offline_trt_8b.jsonl +cat offline_trt_8b.jsonl | cut -d':' -f12 | cut -d',' -f1 +``` + +```bash +# TensorRT LLM Online 8B + +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online_trt_8b.jsonl +cat online_trt_8b.jsonl | cut -d':' -f9 | cut -d',' -f1 +``` + +```bash +# TensorRT LLM Offline 70B + +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline_trt_70b.jsonl +python3 bench_serving.py --backend trt --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline_trt_70b.jsonl --model meta-llama/Meta-Llama-3-70B-Instruct +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name sharegpt --num-prompts 3000 --output-file offline_trt_70b.jsonl +cat offline_trt_70b.jsonl | cut -d':' -f12 | cut -d',' -f1 +``` + +```bash +# TensorRT LLM Online 70B + +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online_trt_70b.jsonl +cat online_trt_70b.jsonl | cut -d':' -f9 | cut -d',' -f1 +``` diff --git a/benchmark/blog_v0_2/config.md b/benchmark/blog_v0_2/config.md new file mode 100644 index 000000000..3faf6009b --- /dev/null +++ b/benchmark/blog_v0_2/config.md @@ -0,0 +1,100 @@ +### used for TensorRT LLM + +``` +{ + "architecture": "LlamaForCausalLM", + "dtype": "float16", + "logits_dtype": "float32", + "vocab_size": 128256, + "max_position_embeddings": 8192, + "hidden_size": 16384, + "num_hidden_layers": 126, + "num_attention_heads": 128, + "num_key_value_heads": 16, + "head_size": 128, + "qk_layernorm": false, + "hidden_act": "silu", + "intermediate_size": 53248, + "norm_epsilon": 1e-05, + "position_embedding_type": "rope_gpt_neox", + "use_parallel_embedding": false, + "embedding_sharding_dim": 0, + "share_embedding_table": false, + "mapping": { + "world_size": 8, + "tp_size": 8, + "pp_size": 1, + "gpus_per_node": 8 + }, + "quantization": { + "quant_algo": "FP8", + "kv_cache_quant_algo": null, + "group_size": 128, + "smoothquant_val": null, + "has_zero_point": false, + "pre_quant_scale": false, + "exclude_modules": [ + "lm_head" + ] + }, + "kv_dtype": "float16", + "rotary_scaling": null, + "residual_mlp": false, + "moe_normalization_mode": null, + "rotary_base": 500000.0, + "moe_num_experts": 0, + "moe_top_k": 0, + "moe_tp_mode": 2, + "attn_bias": false, + "disable_weight_only_quant_plugin": false, + "mlp_bias": false +} +``` + +### used for vLLM and SGLang + +``` +{ + "_name_or_path": "dummy_fp8", + "architectures": [ + "LlamaForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 128000, + "eos_token_id": 128009, + "hidden_act": "silu", + "hidden_size": 16384, + "initializer_range": 0.02, + "intermediate_size": 53248, + "mlp_bias": false, + "model_type": "llama", + "num_attention_heads": 128, + "num_hidden_layers": 126, + "num_key_value_heads": 8, + "pretraining_tp": 1, + "quantization_config": { + "activation_scheme": "static", + "ignored_layers": [ + "lm_head" + ], + "quant_method": "fp8" + }, + "rope_scaling": { + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3" + }, + "max_position_embeddings": 131072, + "rms_norm_eps": 1e-05, + "rope_scaling": null, + "rope_theta": 500000.0, + "tie_word_embeddings": false, + "torch_dtype": "bfloat16", + "transformers_version": "4.41.1", + "use_cache": true, + "vocab_size": 128256 +} +``` diff --git a/benchmark/boolq/README.md b/benchmark/boolq/README.md new file mode 100644 index 000000000..3704742ee --- /dev/null +++ b/benchmark/boolq/README.md @@ -0,0 +1,19 @@ +## Download data +``` +git clone https://hf-mirror.com/datasets/google/boolq +``` + +## Convert parquet to json +``` +bash parquet_to_json.sh +``` +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path ramblingpolymath/Qwen3-32B-W8A8 --port 30000 +``` + +``` +python3 bench_sglang.py +``` diff --git a/benchmark/boolq/bench_sglang.py b/benchmark/boolq/bench_sglang.py new file mode 100644 index 000000000..b3ce3c996 --- /dev/null +++ b/benchmark/boolq/bench_sglang.py @@ -0,0 +1,124 @@ +import argparse +import json +import time + +import numpy as np + +from sglang.api import set_default_backend +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import read_jsonl + + +def get_example(lines, i, answer): + prompt = "Question: " + lines[i]["question"] + lines[i]["passage"] + "\nAnswer:" + if answer: + prompt += str(lines[i]["answer"]) + return prompt + + +def few_shot_examples(lines, k): + prompts = "" + for i in range(k): + prompts += get_example(lines, i, True) + "\n\n" + return prompts + + +def main(args): + # Select backend + set_default_backend(select_sglang_backend(args)) + + # Read data + train_data_path = args.train_data_path + test_data_path = args.test_data_path + lines_train = list(read_jsonl(train_data_path)) + lines_test = list(read_jsonl(test_data_path)) + + # Construct prompts + num_questions = args.num_questions + num_shots = args.num_shots + few_shots = few_shot_examples(lines_train, num_shots) + + questions = [] + answer = [] + for i in range(len(lines_test[:num_questions])): + questions.append(get_example(lines_test, i, False)) + answer.append(str(lines_test[i]["answer"])) + arguments = [{"question": q} for q in questions] + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + @sgl.function + def few_shot_boolq(s, question): + s += few_shots + question + s += sgl.gen("answer", max_tokens=5, stop=["\n"]) + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Run requests + tic = time.perf_counter() + states = few_shot_boolq.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + preds = [] + for i in range(len(states)): + preds.append(states[i]["answer"]) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(answer)) + + # Compute speed + num_output_tokens = sum( + s.get_meta_info("answer")["completion_tokens"] for s in states + ) + output_throughput = num_output_tokens / latency + + # Print results + print(f"Accuracy: {acc:.3f}") + print(f"Latency: {latency:.3f} s") + print(f"Output throughput: {output_throughput:.3f} token/s") + + # Results + with open(args.result_file, "a") as fout: + value = { + "task": "boolq", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shots", type=int, default=5) + parser.add_argument( + "--train-data-path", type=str, default="./boolq/data/train-00000-of-00001.json" + ) + parser.add_argument( + "--test-data-path", + type=str, + default="./boolq/data/validation-00000-of-00001.json", + ) + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/boolq/convert_parquet_to_json.py b/benchmark/boolq/convert_parquet_to_json.py new file mode 100644 index 000000000..e3e69cb31 --- /dev/null +++ b/benchmark/boolq/convert_parquet_to_json.py @@ -0,0 +1,28 @@ +import sys + +import pyarrow.parquet as pq + + +def convert_parquet_to_json(input_file, output_file): + # read parquet file + table = pq.read_table(input_file) + + # turn parquet data to dataframe + df = table.to_pandas() + + # turn dataframe to json form + json_data = df.to_json(orient="records", lines=True) + + # write json to file + with open(output_file, "w") as f: + f.write(json_data) + + +if __name__ == "__main__": + if len(sys.argv) != 3: + print("Usage:python convert_parquet_to_json.py ") + + input_file = sys.argv[1] + output_file = sys.argv[2] + + convert_parquet_to_json(input_file, output_file) diff --git a/benchmark/boolq/parquet_to_json.sh b/benchmark/boolq/parquet_to_json.sh new file mode 100755 index 000000000..9aaf087ff --- /dev/null +++ b/benchmark/boolq/parquet_to_json.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +#define input and output direction +input_dir="./boolq/data" +output_dir="./boolq/data" + +#define files needed to be handled +files=( + "train-00000-of-00001.parquet" + "validation-00000-of-00001.parquet" +) + +#foe files above, use python script to convert the form +for file in "${files[@]}"; do + input_file="${input_dir}/${file}" + output_file="${output_dir}/${file%.parquet}.json" + + echo "Converting ${input_file} to ${output_file} ..." + python3 convert_parquet_to_json.py "${input_file}" "${output_file}" + + if [ $? -eq 0 ]; then + echo "Conversion successful: ${output_file}" + else + echo "Conversion failed: ${input_file}" + fi +done diff --git a/benchmark/ceval/README.md b/benchmark/ceval/README.md new file mode 100644 index 000000000..b822e43c3 --- /dev/null +++ b/benchmark/ceval/README.md @@ -0,0 +1,15 @@ +## Download data +``` +git lfs clone https://huggingface.co/datasets/ceval/ceval-exam +``` + +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path ramblingpolymath/Qwen3-32B-W8A8 --port 30000 +``` + +``` +python3 bench_sglang.py +``` diff --git a/benchmark/ceval/bench_sglang.py b/benchmark/ceval/bench_sglang.py new file mode 100644 index 000000000..32ed0baf2 --- /dev/null +++ b/benchmark/ceval/bench_sglang.py @@ -0,0 +1,138 @@ +import argparse +import json +import os +import random +import re +import time + +import numpy as np +from datasets import load_dataset + +from sglang.api import set_default_backend +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) + +choices = ["A", "B", "C", "D"] + + +def get_one_example(line, include_answer): + res = line["question"] + res += f"\nA. {line['A']}" + res += f"\nB. {line['B']}" + res += f"\nC. {line['C']}" + res += f"\nD. {line['D']}" + + if include_answer: + res += f"\nAnswer: {line['answer']} \n\n" + return res + + +def get_few_shot_examples(lines): + res = "" + for line in lines: + res += get_one_example(line, True) + "\n\n" + return res + + +def get_answer_value(response): + pattern = r"(Answer:|answer:|答案是|答案是:|正确答案是:|答案:|Assistant:)\s*([A-D])(?![\w])" + match = re.search(pattern, response) + + if match: + return match.group(2) + + return random.choice(choices) + + +def main(args): + # Read data && Construct prompts + arguments = [] + labels = [] + examples = "examples:\n" + data_path = args.data_path + for subject in os.listdir(data_path): + subject_path = os.path.join(data_path, subject) + if os.path.isdir(subject_path) and subject != ".git": + dataset = load_dataset(data_path, name=subject) + dev_lines_temp = dataset["dev"] + val_lines_temp = dataset["val"] + few_shot_examples = get_few_shot_examples(dev_lines_temp, subject) + examples += f"{few_shot_examples}" + for val_line in val_lines_temp: + arguments.append( + { + "examples": few_shot_examples, + "question": get_one_example(val_line, False), + } + ) + labels.append(val_line["answer"]) + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + @sgl.function + def few_shot_ceval(s, examples, question): + s += examples + question + sgl.gen("Answer") + + ##################################### + ########## SGL Program End ########## + ##################################### + + num_questions = args.num_questions if args.num_questions else len(arguments) + + # Select backend + set_default_backend(select_sglang_backend(args)) + + # Run requests + tic = time.perf_counter() + states = few_shot_ceval.run_batch( + arguments[:num_questions], + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + preds = [get_answer_value(states[i]["Answer"]) for i in range(num_questions)] + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels[:num_questions])) + + # Compute speed + num_output_tokens = sum( + s.get_meta_info("Answer")["completion_tokens"] for s in states + ) + output_throughput = num_output_tokens / latency + + # Print results + print(f"Accuracy: {acc:.3f}") + print(f"Latency: {latency:.3f} s") + print(f"Output throughput: {output_throughput:.3f} token/s") + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "ceval", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="ceval-exam") + parser.add_argument("--num-questions", type=int, default=None) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/deepseek_v3/README.md b/benchmark/deepseek_v3/README.md new file mode 100644 index 000000000..53efc23f5 --- /dev/null +++ b/benchmark/deepseek_v3/README.md @@ -0,0 +1,373 @@ +# DeepSeek V3.1/V3/R1 Support + +The SGLang and DeepSeek teams collaborated to get DeepSeek V3 FP8 running on NVIDIA and AMD GPUs **from day one**. SGLang also supports [MLA optimization](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#deepseek-multi-head-latent-attention-mla-throughput-optimizations) and [DP attention](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models), making SGLang one of the best open-source LLM engines for running DeepSeek models. SGLang is the inference engine recommended by the official [DeepSeek team](https://github.com/deepseek-ai/DeepSeek-V3/tree/main?tab=readme-ov-file#62-inference-with-sglang-recommended). + +Special thanks to Meituan's Search & Recommend Platform Team and Baseten's Model Performance Team for implementing the model, and DataCrunch for providing GPU resources. + +For optimizations made on the DeepSeek series models regarding SGLang, please refer to [DeepSeek Model Optimizations in SGLang](https://docs.sglang.ai/basic_usage/deepseek.html). + +## Installation & Launch + +If you encounter errors when starting the server, ensure the weights have finished downloading. It's recommended to download them beforehand or restart multiple times until all weights are downloaded. + +### Using Docker (Recommended) + +```bash +# Pull latest image +# https://hub.docker.com/r/lmsysorg/sglang/tags +docker pull lmsysorg/sglang:latest + +# Launch +docker run --gpus all --shm-size 32g -p 30000:30000 -v ~/.cache/huggingface:/root/.cache/huggingface --ipc=host --network=host --privileged lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code --port 30000 +``` + +If you are using RDMA, please note that: + +1. `--network host` and `--privileged` are required by RDMA. If you don't need RDMA, you can remove them. +2. You may need to set `NCCL_IB_GID_INDEX` if you are using RoCE, for example: `export NCCL_IB_GID_INDEX=3`. + +Add [performance optimization options](#performance-optimization-options) as needed. + +### Using pip + +```bash +# Installation +pip install "sglang[all]>=0.5.2" + +# Launch +python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code +``` + +Add [performance optimization options](#performance-optimization-options) as needed. + + + +### Performance Optimization Options + +[MLA optimizations](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#deepseek-multi-head-latent-attention-mla-throughput-optimizations) are enabled by default. Here are some optional optimizations can be enabled as needed. + +- [Data Parallelism Attention](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models): For high QPS scenarios, add the `--enable-dp-attention` argument to boost throughput. +- [Torch.compile Optimization](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#torchcompile-latency-optimizations): Add `--enable-torch-compile` argument to enable it. This will take some time while server starts. The maximum batch size for torch.compile optimization can be controlled with `--torch-compile-max-bs`. It's recommended to set it between `1` and `8`. (e.g., `--torch-compile-max-bs 8`) + +### Usage: Chat with DeepSeek + +#### DeepSeek V3/R1 + +```python3 +import openai +client = openai.Client( + base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") + +# Chat completion +response = client.chat.completions.create( + model="default", + messages=[ + {"role": "system", "content": "You are a helpful AI assistant"}, + {"role": "user", "content": "List 3 countries and their capitals."}, + ], + temperature=0, + max_tokens=64, +) +print(response) +``` + +#### DeepSeek V3.1 +On top of the basic usage similar to the DeepSeek V3/R1 example, DeepSeek V3.1 supports a request-level thinking/non-thinking toggle. Simply switch the `"thinking"` field in `extra_body={"chat_template_kwargs": {"thinking": True}}` to enable/disable the thinking mode. + +##### Non Thinking +```python3 +import openai +client = openai.Client( + base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") + +# Chat completion +response = client.chat.completions.create( + model="default", + messages=[ + {"role": "system", "content": "You are a helpful AI assistant"}, + {"role": "user", "content": "Answer the following with the second letter of the correct answer only: What is the capital of France?"}, + ], + temperature=0, + max_tokens=1024, + extra_body = {"chat_template_kwargs": {"thinking": False}} +) +print(response.choices[0].message.content) +``` +Answer: +``` +h +``` +* The correct response should be 'A', as the correct answer to the question is 'Paris'. +##### Thinking +```python3 +import openai +client = openai.Client( + base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") + +# Chat completion +response = client.chat.completions.create( + model="default", + messages=[ + {"role": "system", "content": "You are a helpful AI assistant"}, + {"role": "user", "content": "Answer the following with the second letter of the correct answer only: What is the capital of France?"}, + ], + temperature=0, + max_tokens=1024, + extra_body = {"chat_template_kwargs": {"thinking": True}} +) +print(response) +``` +Answer: +``` +First, the question is: "What is the capital of France?" I know that the capital of France is Paris. + +The user says: "Answer the following with the second letter of the correct answer only." So, I need to provide only the second letter of the correct answer. + +The correct answer is "Paris". Now, I need to find the second letter of "Paris". + +Let's spell it out: P-A-R-I-S. + +- First letter: P + +- Second letter: A + +- Third letter: R + +- Fourth letter: I + +- Fifth letter: S + +So, the second letter is "A". + +I should only output the second letter, which is "A". No additional text or explanation, just the letter. + +The user emphasized "the second letter of the correct answer only", so my response should be just "A". + +Finally, I need to make sure that this is the correct answer. Yes, Paris is indeed the capital of France.A +``` +* The response contains `` thinking trace and model was able to derive the correct answer from it. + +### Example: Serving with two H20\*8 nodes + +For example, there are two H20 nodes, each with 8 GPUs. The first node's IP is `10.0.0.1`, and the second node's IP is `10.0.0.2`. Please **use the first node's IP** for both commands. + +If the command fails, try setting the `GLOO_SOCKET_IFNAME` parameter. For more information, see [Common Environment Variables](https://pytorch.org/docs/stable/distributed.html#common-environment-variables). + +If the multi nodes support NVIDIA InfiniBand and encounter hanging issues during startup, consider adding the parameter `export NCCL_IB_GID_INDEX=3`. For more information, see [this](https://github.com/sgl-project/sglang/issues/3516#issuecomment-2668493307). + +```bash +# node 1 +python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 10.0.0.1:5000 --nnodes 2 --node-rank 0 --trust-remote-code + +# node 2 +python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 10.0.0.1:5000 --nnodes 2 --node-rank 1 --trust-remote-code +``` + +If you have two H100 nodes, the usage is similar to the aforementioned H20. + +> **Note that the launch command here does not enable Data Parallelism Attention or `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). + +### Example: Serving with two H200\*8 nodes and docker + +There are two H200 nodes, each with 8 GPUs. The first node's IP is `192.168.114.10`, and the second node's IP is `192.168.114.11`. Configure the endpoint to expose it to another Docker container using `--host 0.0.0.0` and `--port 40000`, and set up communications with `--dist-init-addr 192.168.114.10:20000`. +A single H200 with 8 devices can run DeepSeek V3, the dual H200 setup is just to demonstrate multi-node usage. + +```bash +# node 1 +docker run --gpus all \ + --shm-size 32g \ + --network=host \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --name sglang_multinode1 \ + -it \ + --rm \ + --env "HF_TOKEN=$HF_TOKEN" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 192.168.114.10:20000 --nnodes 2 --node-rank 0 --trust-remote-code --host 0.0.0.0 --port 40000 +``` + +```bash +# node 2 +docker run --gpus all \ + --shm-size 32g \ + --network=host \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --name sglang_multinode2 \ + -it \ + --rm \ + --env "HF_TOKEN=$HF_TOKEN" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 192.168.114.10:20000 --nnodes 2 --node-rank 1 --trust-remote-code --host 0.0.0.0 --port 40000 +``` + +To ensure functionality, we include a test from a client Docker container. + +```bash +docker run --gpus all \ + --shm-size 32g \ + --network=host \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --name sglang_multinode_client \ + -it \ + --rm \ + --env "HF_TOKEN=$HF_TOKEN" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1 --random-output 512 --random-range-ratio 1 --num-prompts 1 --host 0.0.0.0 --port 40000 --output-file "deepseekv3_multinode.jsonl" +``` + +> **Note that the launch command here does not enable Data Parallelism Attention or `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). + +### Example: Serving with four A100\*8 nodes + +To serve DeepSeek-V3 with A100 GPUs, we need to convert the [FP8 model checkpoints](https://huggingface.co/deepseek-ai/DeepSeek-V3) to BF16 with [script](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/fp8_cast_bf16.py) mentioned [here](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/fp8_cast_bf16.py) first. + +Since the BF16 model is over 1.3 TB, we need to prepare four A100 nodes, each with 8 80GB GPUs. Assume the first node's IP is `10.0.0.1`, and the converted model path is `/path/to/DeepSeek-V3-BF16`, we can have following commands to launch the server. + +```bash +# node 1 +python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 0 --trust-remote-code --host 0.0.0.0 --port 30000 + +# node 2 +python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 1 --trust-remote-code + +# node 3 +python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 2 --trust-remote-code + +# node 4 +python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 3 --trust-remote-code +``` + +> **Note that the launch command here does not enable Data Parallelism Attention or `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). + +Then we can benchmark the accuracy and latency by accessing the first node's exposed port with the following example commands. + +```bash +# bench accuracy +python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --host http://10.0.0.1 --port 30000 + +# bench latency +python3 -m sglang.bench_one_batch_server --model None --base-url http://10.0.0.1:30000 --batch-size 1 --input-len 128 --output-len 128 +``` + + +### Example: Serving with 8 A100/A800 with AWQ Quantization + +**Recommended Usage** + +Add `--quantization moe_wna16` flag to enable moe wna16 kernel for better performance. +One example is as follows: + +```bash +python3 -m sglang.launch_server --model cognitivecomputations/DeepSeek-R1-AWQ --tp 8 --trust-remote-code --quantization moe_wna16 +``` + +Alternatively, you can use `--quantization awq_marlin` as follows: + +```bash +python3 -m sglang.launch_server --model cognitivecomputations/DeepSeek-R1-AWQ --tp 8 --trust-remote-code --quantization awq_marlin --dtype float16 +``` + +Note that `awq_marlin` only supports `float16` now, which may lead to some precision loss. + +### Example: Serving with 16 A100/A800 with int8 Quantization + +There are block-wise and per-channel quantization methods, and the quantization parameters have already been uploaded to Huggingface. One example is as follows: + +- [meituan/DeepSeek-R1-Block-INT8](https://huggingface.co/meituan/DeepSeek-R1-Block-INT8) +- [meituan/DeepSeek-R1-Channel-INT8](https://huggingface.co/meituan/DeepSeek-R1-Channel-INT8) + +Assuming that master node IP is `MASTER_IP`, checkpoint path is `/path/to/DeepSeek-R1-INT8` and port=5000, we can have following commands to launch the server: +```bash +#master +python3 -m sglang.launch_server \ + --model meituan/DeepSeek-R1-Block-INT8 --tp 16 --dist-init-addr \ + MASTER_IP:5000 --nnodes 2 --node-rank 0 --trust-remote-code --enable-torch-compile --torch-compile-max-bs 8 +#cluster +python3 -m sglang.launch_server \ + --model meituan/DeepSeek-R1-Block-INT8 --tp 16 --dist-init-addr \ + MASTER_IP:5000 --nnodes 2 --node-rank 1 --trust-remote-code --enable-torch-compile --torch-compile-max-bs 8 +``` + +> **Note that the launch command here enables `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). + +Then on the **master node**, supposing the ShareGPT data is located at `/path/to/ShareGPT_V3_unfiltered_cleaned_split.json`, you can run the following commands to benchmark the launched server: + +```bash +# bench accuracy +python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 + +# bench serving +python3 -m sglang.bench_serving --dataset-path /path/to/ShareGPT_V3_unfiltered_cleaned_split.json --dataset-name random --random-input 128 --random-output 128 --num-prompts 1000 --request-rate 128 --random-range-ratio 1.0 +``` + +> **Note: using `--parallel 200` can accelerate accuracy benchmarking**. + +### Example: Serving with 32 L40S with int8 Quantization + +Running with per-channel quantization model: + +- [meituan/DeepSeek-R1-Channel-INT8](https://huggingface.co/meituan/DeepSeek-R1-Channel-INT8) + +Assuming that master node IP is `MASTER_IP`, checkpoint path is `/path/to/DeepSeek-R1-Channel-INT8` and port=5000, we can have following commands to launch the server: + +```bash +#master +python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ + --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 0 --trust-remote \ + --enable-torch-compile --torch-compile-max-bs 32 +#cluster +python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ + --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 1 --trust-remote \ + --enable-torch-compile --torch-compile-max-bs 32 +python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ + --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 2 --trust-remote \ + --enable-torch-compile --torch-compile-max-bs 32 +python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ + --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 3 --trust-remote \ + --enable-torch-compile --torch-compile-max-bs 32 +``` + +The benchmarking method is the same as describted in the previous [16 x A100](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-16-a100a800-with-int8-quantization) example. + +### Example: Serving on any cloud or Kubernetes with SkyPilot + +SkyPilot helps find cheapest available GPUs across any cloud or existing Kubernetes clusters and launch distributed serving with a single command. See details [here](https://github.com/skypilot-org/skypilot/tree/master/llm/deepseek-r1). + +To serve on multiple nodes: + +```bash +git clone https://github.com/skypilot-org/skypilot.git +# Serve on 2 H100/H200x8 nodes +sky launch -c r1 llm/deepseek-r1/deepseek-r1-671B.yaml --retry-until-up +# Serve on 4 A100x8 nodes +sky launch -c r1 llm/deepseek-r1/deepseek-r1-671B-A100.yaml --retry-until-up +``` + +#### Troubleshooting + +If you encounter the following error with fp16/bf16 checkpoint: + +```bash +ValueError: Weight output_partition_size = 576 is not divisible by weight quantization block_n = 128. +``` + +edit your `config.json` and remove the `quantization_config` block. For example: + +```json +"quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "weight_block_size": [128, 128] +}, +``` + +Removing this block typically resolves the error. For more details, see the discussion in [sgl-project/sglang#3491](https://github.com/sgl-project/sglang/issues/3491#issuecomment-2650779851). + +## DeepSeek V3 Optimization Plan + +https://github.com/sgl-project/sglang/issues/2591 diff --git a/benchmark/dspy/README.md b/benchmark/dspy/README.md new file mode 100644 index 000000000..a2b213aa6 --- /dev/null +++ b/benchmark/dspy/README.md @@ -0,0 +1,51 @@ +## Install + +``` +pip3 install dspy-ai +``` + +Turn off cache at https://github.com/stanfordnlp/dspy/blob/34d8420383ec752037aa271825c1d3bf391e1277/dsp/modules/cache_utils.py#L10. +``` +cache_turn_on = False +``` + +or set the environment variable + +``` +export DSP_CACHEBOOL=false +``` + +## Benchmark SGLang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_dspy_intro.py --backend sglang +``` + + +## Benchmark TGI +``` +docker run --name tgi --rm -ti --gpus all --network host \ + -v /home/ubuntu/model_weights/Llama-2-7b-chat-hf:/Llama-2-7b-chat-hf \ + ghcr.io/huggingface/text-generation-inference:1.3.0 \ + --model-id /Llama-2-7b-chat-hf --num-shard 1 --trust-remote-code \ + --max-input-length 2048 --max-total-tokens 4096 \ + --port 24000 +``` + +``` +python3 bench_dspy_intro.py --backend tgi +``` + + + +## Benchmark vLLM +``` +python3 -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_dspy_intro.py --backend vllm +``` diff --git a/benchmark/dspy/bench_dspy_intro.py b/benchmark/dspy/bench_dspy_intro.py new file mode 100644 index 000000000..2b0936ed2 --- /dev/null +++ b/benchmark/dspy/bench_dspy_intro.py @@ -0,0 +1,192 @@ +""" +Adapted from +https://github.com/stanfordnlp/dspy/blob/34d8420383ec752037aa271825c1d3bf391e1277/intro.ipynb#L9 +""" + +import argparse + +import dspy +from dspy.datasets import HotPotQA + + +class BasicQA(dspy.Signature): + """Answer questions with short factoid answers.""" + + question = dspy.InputField() + answer = dspy.OutputField(desc="often between 1 and 5 words") + + +class GenerateAnswer(dspy.Signature): + """Answer questions with short factoid answers.""" + + context = dspy.InputField(desc="may contain relevant facts") + question = dspy.InputField() + answer = dspy.OutputField(desc="often between 1 and 5 words") + + +class RAG(dspy.Module): + def __init__(self, num_passages=3): + super().__init__() + + self.retrieve = dspy.Retrieve(k=num_passages) + self.generate_answer = dspy.ChainOfThought(GenerateAnswer) + + def forward(self, question): + context = self.retrieve(question).passages + prediction = self.generate_answer(context=context, question=question) + return dspy.Prediction(context=context, answer=prediction.answer) + + +def main(args): + # lm = dspy.OpenAI(model='gpt-3.5-turbo') + if args.backend == "tgi": + lm = dspy.HFClientTGI( + model="meta-llama/Llama-2-7b-chat-hf", + port=args.port, + url="http://localhost", + ) + elif args.backend == "sglang": + lm = dspy.HFClientSGLang( + model="meta-llama/Llama-2-7b-chat-hf", + port=args.port, + url="http://localhost", + ) + elif args.backend == "vllm": + lm = dspy.HFClientVLLM( + model="meta-llama/Llama-2-7b-chat-hf", + port=args.port, + url="http://localhost", + ) + else: + raise ValueError(f"Invalid backend: {args.backend}") + + colbertv2_wiki17_abstracts = dspy.ColBERTv2( + url="http://20.102.90.50:2017/wiki17_abstracts" + ) + dspy.settings.configure(lm=lm, rm=colbertv2_wiki17_abstracts) + + # Load the dataset. + dataset = HotPotQA( + train_seed=1, train_size=20, eval_seed=2023, dev_size=args.dev_size, test_size=0 + ) + + # Tell DSPy that the 'question' field is the input. Any other fields are labels and/or metadata. + trainset = [x.with_inputs("question") for x in dataset.train] + devset = [x.with_inputs("question") for x in dataset.dev] + + print(len(trainset), len(devset)) + + train_example = trainset[0] + print(f"Question: {train_example.question}") + print(f"Answer: {train_example.answer}") + + dev_example = devset[18] + print(f"Question: {dev_example.question}") + print(f"Answer: {dev_example.answer}") + print(f"Relevant Wikipedia Titles: {dev_example.gold_titles}") + + print( + f"For this dataset, training examples have input keys {train_example.inputs().keys()} and label keys {train_example.labels().keys()}" + ) + print( + f"For this dataset, dev examples have input keys {dev_example.inputs().keys()} and label keys {dev_example.labels().keys()}" + ) + + # Define the predictor. + generate_answer = dspy.Predict(BasicQA) + + # Call the predictor on a particular input. + pred = generate_answer(question=dev_example.question) + + # Print the input and the prediction. + print(f"Question: {dev_example.question}") + print(f"Predicted Answer: {pred.answer}") + + lm.inspect_history(n=1) + + # Define the predictor. Notice we're just changing the class. The signature BasicQA is unchanged. + generate_answer_with_chain_of_thought = dspy.ChainOfThought(BasicQA) + + # Call the predictor on the same input. + pred = generate_answer_with_chain_of_thought(question=dev_example.question) + + # Print the input, the chain of thought, and the prediction. + print(f"Question: {dev_example.question}") + print(f"Thought: {pred.rationale.split('.', 1)[1].strip()}") + print(f"Predicted Answer: {pred.answer}") + + retrieve = dspy.Retrieve(k=3) + topK_passages = retrieve(dev_example.question).passages + + print( + f"Top {retrieve.k} passages for question: {dev_example.question} \n", + "-" * 30, + "\n", + ) + + for idx, passage in enumerate(topK_passages): + print(f"{idx+1}]", passage, "\n") + + retrieve("When was the first FIFA World Cup held?").passages[0] + + from dspy.teleprompt import BootstrapFewShot + + # Validation logic: check that the predicted answer is correct. + # Also check that the retrieved context does actually contain that answer. + def validate_context_and_answer(example, pred, trace=None): + answer_EM = dspy.evaluate.answer_exact_match(example, pred) + answer_PM = dspy.evaluate.answer_passage_match(example, pred) + return answer_EM and answer_PM + + # Set up a basic teleprompter, which will compile our RAG program. + teleprompter = BootstrapFewShot(metric=validate_context_and_answer) + + # Compile! + compiled_rag = teleprompter.compile(RAG(), trainset=trainset) + + # Ask any question you like to this simple RAG program. + my_question = "What castle did David Gregory inherit?" + + # Get the prediction. This contains `pred.context` and `pred.answer`. + pred = compiled_rag(my_question) + + # Print the contexts and the answer. + print(f"Question: {my_question}") + print(f"Predicted Answer: {pred.answer}") + print(f"Retrieved Contexts (truncated): {[c[:200] + '...' for c in pred.context]}") + + from dspy.evaluate.evaluate import Evaluate + + # Set up the `evaluate_on_hotpotqa` function. We'll use this many times below. + evaluate_on_hotpotqa = Evaluate( + devset=devset, + num_threads=args.num_threads, + display_progress=True, + display_table=5, + ) + + # Evaluate the `compiled_rag` program with the `answer_exact_match` metric. + metric = dspy.evaluate.answer_exact_match + evaluate_on_hotpotqa(compiled_rag, metric=metric) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--port", type=int) + parser.add_argument("--num-threads", type=int, default=32) + parser.add_argument("--dev-size", type=int, default=150) + parser.add_argument( + "--backend", type=str, choices=["sglang", "tgi", "vllm"], default="sglang" + ) + args = parser.parse_args() + + if args.port is None: + default_port = { + "vllm": 21000, + "lightllm": 22000, + "tgi": 24000, + "sglang": 30000, + } + args.port = default_port.get(args.backend, None) + + main(args) diff --git a/benchmark/generative_agents/README.md b/benchmark/generative_agents/README.md new file mode 100644 index 000000000..393a9ce83 --- /dev/null +++ b/benchmark/generative_agents/README.md @@ -0,0 +1,38 @@ +## Download the dataset + +``` +wget -O agent_calls.jsonl https://drive.google.com/uc?export=download&id=19qLpD45e9JGTKF2cUjJJegwzSUEZEKht +``` + +## Run benchmark + +Ensure that this benchmark is run in a serial manner (using --parallel 1) to preserve any potential dependencies between requests. + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-events 1000 --parallel 1 +``` + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-events 1000 --backend vllm --parallel 1 +``` + +### Benchmark guidance +``` +python3 bench_other.py --num-events 1000 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --num-events 1000 --backend lmql --parallel 1 +``` diff --git a/benchmark/generative_agents/agent_functions.py b/benchmark/generative_agents/agent_functions.py new file mode 100644 index 000000000..9785fc1ec --- /dev/null +++ b/benchmark/generative_agents/agent_functions.py @@ -0,0 +1,300 @@ +import sglang as sgl + +# here are the top five agent functions contributing ~70% LLM calls +# reference: https://github.com/joonspk-research/generative_agents/ + + +@sgl.function +def poignancy_event(s, persona_name, persona_iss, event): + s += "Here is a brief description of " + persona_name + ".\n" + s += persona_iss + "\n" + s += "On the scale of 1 to 10, where 1 is purely mundane (e.g., brushing teeth, making bed) and 10 is extremely poignant (e.g., a break up, college acceptance), rate the likely poignancy of the following event for" + s += persona_name + ".\n\n" + s += "Event: " + event + s += "Rate (return a number between 1 to 10):" + s += sgl.gen(name="Rate", max_tokens=2) + + +def poignancy_event_prompt(persona_name, persona_iss, event): + # return prompt and max_tokens + s = "" + s += "Here is a brief description of " + persona_name + ".\n" + s += persona_iss + "\n" + s += "On the scale of 1 to 10, where 1 is purely mundane (e.g., brushing teeth, making bed) and 10 is extremely poignant (e.g., a break up, college acceptance), rate the likely poignancy of the following event for" + s += persona_name + ".\n\n" + s += "Event: " + event + s += "Rate (return a number between 1 to 10):" + return {"prompt": s, "max_tokens": 2, "stop": None} + + +@sgl.function +def generate_event_triple(s, persona_name, action): + s += """Task: Turn the input into (subject, predicate, object). +Input: Sam Johnson is eating breakfast. +Output: (Dolores Murphy, eat, breakfast) +--- +Input: Joon Park is brewing coffee. +Output: (Joon Park, brew, coffee) +--- +Input: Jane Cook is sleeping. +Output: (Jane Cook, is, sleep) +--- +Input: Michael Bernstein is writing email on a computer. +Output: (Michael Bernstein, write, email) +--- +Input: Percy Liang is teaching students in a classroom. +Output: (Percy Liang, teach, students) +--- +Input: Merrie Morris is running on a treadmill. +Output: (Merrie Morris, run, treadmill) +---""" + s += persona_name + "is" + action + ".\n" + s += "(" + persona_name + "," + s += sgl.gen(name="Triple", max_tokens=20, stop=")") + + +def generate_event_triple_prompt(persona_name, action): + s = "" + s += """Task: Turn the input into (subject, predicate, object). +Input: Sam Johnson is eating breakfast. +Output: (Dolores Murphy, eat, breakfast) +--- +Input: Joon Park is brewing coffee. +Output: (Joon Park, brew, coffee) +--- +Input: Jane Cook is sleeping. +Output: (Jane Cook, is, sleep) +--- +Input: Michael Bernstein is writing email on a computer. +Output: (Michael Bernstein, write, email) +--- +Input: Percy Liang is teaching students in a classroom. +Output: (Percy Liang, teach, students) +--- +Input: Merrie Morris is running on a treadmill. +Output: (Merrie Morris, run, treadmill) +---""" + s += persona_name + "is" + action + ".\n" + s += "(" + persona_name + "," + return {"prompt": s, "max_tokens": 20, "stop": ")"} + + +@sgl.function +def generate_pronunciatio(s, action): + s += "Convert an action description to an emoji (important: use two or less emojis).\n" + s += "Action description: " + action + ".\n" + s += "Emoji:" + sgl.gen(name="Emoji", max_tokens=6) + + +def generate_pronunciatio_prompt(action): + s = "" + s += "Convert an action description to an emoji (important: use two or less emojis).\n" + s += "Action description: " + action + ".\n" + s += "Emoji:" + return {"prompt": s, "max_tokens": 6, "stop": None} + + +@sgl.function +def action_location_sector( + s, + persona_name, + living_sector, + living_sector_areas, + current_sector, + current_sector_areas, + daily_plan, + sector_options, + current_action, + next_action, +): + s += """Task -- choose an appropriate area from the area options for a task at hand. +Sam Kim lives in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. +Sam Kim is currently in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. +Area options: {Sam Kim's house, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. +* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim. +For taking a walk, Sam Kim should go to the following area: {Johnson Park} +--- +Jane Anderson lives in {Oak Hill College Student Dormatory} that has Jane Anderson's room. +Jane Anderson is currently in {Oak Hill College} that has a classroom, library +Area options: {Oak Hill College Student Dormatory, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. +* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim. +For eating dinner, Jane Anderson should go to the following area: {Hobbs Cafe} +---""" + s += ( + persona_name + + " lives in " + + living_sector + + " that has " + + living_sector_areas + + ".\n" + ) + s += ( + persona_name + + " is currently in " + + current_sector + + " that has " + + current_sector_areas + + ".\n" + ) + s += daily_plan + ".\n" + s += "Area options: " + sector_options + ".\n" + s += """* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim.\n""" + s += ( + persona_name + + " is " + + current_action + + ". For " + + next_action + + ", " + + persona_name + + " should go to the following area: {" + ) + s += sgl.gen(name="Location", max_tokens=10, stop="}") + + +def action_location_sector_prompt( + persona_name, + living_sector, + living_sector_areas, + current_sector, + current_sector_areas, + daily_plan, + sector_options, + current_action, + next_action, +): + s = "" + s += """Task -- choose an appropriate area from the area options for a task at hand. +Sam Kim lives in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. +Sam Kim is currently in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. +Area options: {Sam Kim's house, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. +* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim. +For taking a walk, Sam Kim should go to the following area: {Johnson Park} +--- +Jane Anderson lives in {Oak Hill College Student Dormatory} that has Jane Anderson's room. +Jane Anderson is currently in {Oak Hill College} that has a classroom, library +Area options: {Oak Hill College Student Dormatory, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. +* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim. +For eating dinner, Jane Anderson should go to the following area: {Hobbs Cafe} +---""" + s += ( + persona_name + + " lives in " + + living_sector + + " that has " + + living_sector_areas + + ".\n" + ) + s += ( + persona_name + + " is currently in " + + current_sector + + " that has " + + current_sector_areas + + ".\n" + ) + s += daily_plan + ".\n" + s += "Area options: " + sector_options + ".\n" + s += """* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim.\n""" + s += ( + persona_name + + " is " + + current_action + + ". For " + + next_action + + ", " + + persona_name + + " should go to the following area: {" + ) + return {"prompt": s, "max_tokens": 10, "stop": "}"} + + +@sgl.function +def action_location_object( + s, persona_name, target_sector, target_sector_areas, current_action, next_action +): + s += """ +Jane Anderson is in kitchen in Jane Anderson's house. +Jane Anderson is going to Jane Anderson's house that has the following areas: {kitchen, bedroom, bathroom} +Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. +For cooking, Jane Anderson should go to the following area in Jane Anderson's house: +Answer: {kitchen} +--- +Tom Watson is in common room in Tom Watson's apartment. +Tom Watson is going to Hobbs Cafe that has the following areas: {cafe} +Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. +For getting coffee, Tom Watson should go to the following area in Hobbs Cafe: +Answer: {cafe} +---""" + s += ( + persona_name + + " is going to " + + target_sector + + " that has the following areas: {" + + target_sector_areas + + "}\n" + ) + s += """* Stay in the current area if the activity can be done there. +* NEVER go into other people's rooms unless necessary.""" + s += ( + persona_name + + " is " + + current_action + + ". For " + + next_action + + ", " + + persona_name + + "should go to the following area in " + + target_sector + ) + s += " (MUST pick one of {" + target_sector_areas + "}):\n" + s += "Answer: {" + sgl.gen(name="Area", max_tokens=5, stop="}") + + +def action_location_object_prompt( + persona_name, target_sector, target_sector_areas, current_action, next_action +): + s = "" + s += """ +Jane Anderson is in kitchen in Jane Anderson's house. +Jane Anderson is going to Jane Anderson's house that has the following areas: {kitchen, bedroom, bathroom} +Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. +For cooking, Jane Anderson should go to the following area in Jane Anderson's house: +Answer: {kitchen} +--- +Tom Watson is in common room in Tom Watson's apartment. +Tom Watson is going to Hobbs Cafe that has the following areas: {cafe} +Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. +For getting coffee, Tom Watson should go to the following area in Hobbs Cafe: +Answer: {cafe} +---""" + s += ( + persona_name + + " is going to " + + target_sector + + " that has the following areas: {" + + target_sector_areas + + "}\n" + ) + s += """* Stay in the current area if the activity can be done there. +* NEVER go into other people's rooms unless necessary.""" + s += ( + persona_name + + " is " + + current_action + + ". For " + + next_action + + ", " + + persona_name + + "should go to the following area in " + + target_sector + ) + s += " (MUST pick one of {" + target_sector_areas + "}):\n" + s += "Answer: {" + return {"prompt": s, "max_tokens": 5, "stop": "}"} diff --git a/benchmark/generative_agents/bench_other.py b/benchmark/generative_agents/bench_other.py new file mode 100644 index 000000000..c0b3a3406 --- /dev/null +++ b/benchmark/generative_agents/bench_other.py @@ -0,0 +1,80 @@ +import argparse +import json +import time + +from agent_functions import ( + action_location_object_prompt, + action_location_sector_prompt, + generate_event_triple_prompt, + generate_pronunciatio_prompt, + poignancy_event_prompt, +) +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_events] + mapping = { + "poignancy_event": poignancy_event_prompt, + "generate_event_triple": generate_event_triple_prompt, + "generate_pronunciatio": generate_pronunciatio_prompt, + "action_location_sector": action_location_sector_prompt, + "action_location_object": action_location_object_prompt, + } + + arguments = [mapping[k](**v) for l in lines for k, v in l.items()] + states = [] + + # Select backend + call_generate = get_call_generate(args) + + def get_one_answer(arg): + answer = call_generate(**arg, temperature=0) + states.append(answer) + + async def get_one_answer_async(arg): + answer = await call_generate(**arg, temperature=0) + states.append(answer) + + tic = time.perf_counter() + # we always sequentially execute agent calls to maintain its dependency + if args.backend != "lmql": + for arg in tqdm(arguments): + get_one_answer(arg) + else: + import asyncio + + loop = asyncio.get_event_loop() + for arg in tqdm(arguments): + loop.run_until_complete(get_one_answer_async(arg)) + latency = time.perf_counter() - tic + + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "Generative Agents", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + # to pack weighted functions as a single agent + "num_requests": len(arguments) / len(mapping), + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="agent_calls.jsonl") + parser.add_argument("--num-events", type=int, default=10) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/generative_agents/bench_sglang.py b/benchmark/generative_agents/bench_sglang.py new file mode 100644 index 000000000..034b16591 --- /dev/null +++ b/benchmark/generative_agents/bench_sglang.py @@ -0,0 +1,74 @@ +import argparse +import json +import time + +from agent_functions import ( + action_location_object, + action_location_sector, + generate_event_triple, + generate_pronunciatio, + poignancy_event, +) + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_events] + mapping = { + "poignancy_event": poignancy_event, + "generate_event_triple": generate_event_triple, + "generate_pronunciatio": generate_pronunciatio, + "action_location_sector": action_location_sector, + "action_location_object": action_location_object, + } + arguments = [{mapping[k]: v for k, v in l.items()} for l in lines] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + states = [] + # Run requests + tic = time.perf_counter() + for a in arguments: + # only a single key in the dict + for func, arg in a.items(): + result = func.run(**arg) + result.sync() + states.append(result) + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "Generative Agents", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + # to pack weighted functions as a single agent + "num_requests": len(arguments) / len(mapping), + "other": { + "num_events": args.num_events, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="agent_calls.jsonl") + parser.add_argument("--num-events", type=int, default=10) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/gpt_oss/README.md b/benchmark/gpt_oss/README.md new file mode 100644 index 000000000..4d1b00e91 --- /dev/null +++ b/benchmark/gpt_oss/README.md @@ -0,0 +1,163 @@ +# How to reproduce the result of GPT-OSS with SGLang + +### Install the latest SGLang + +```bash +git clone https://github.com/sgl-project/sglang.git +cd sglang +git checkout v0.5.1.post3 + +pip install --upgrade pip +pip install -e "python[all]" +``` + +### Reproduce the benchmark throughput result (Batch Size 1) + +Launch Command + +```bash +# MXFP4 120B on H100 +python3 -m sglang.launch_server --model openai/gpt-oss-120b --tp 8 --attention-backend triton + +# BF16 120B on H100 +python3 -m sglang.launch_server --model lmsys/gpt-oss-120b-bf16 --tp 8 --attention-backend triton + +# MXFP4 120B on B200 +python3 -m sglang.launch_server --model openai/gpt-oss-120b --tp 4 + +# BF16 120B on B200 +python3 -m sglang.launch_server --model lmsys/gpt-oss-120b-bf16 --tp 4 +``` + +Benchmark Command + +```bash + +# MXFP4 120B on H100 +python3 -m sglang.bench_one_batch_server --model openai/gpt-oss-120b --base-url http://localhost:30000 --batch-size 1 --input-len 1024 --output-len 512 --show-report +``` + +### Reproduce the benchmark throughput result (Batch Size 32) + +Launch Command + +```bash +# MXFP4 120B on H100 +python3 -m sglang.launch_server --model openai/gpt-oss-120b --tp 8 + +# BF16 120B on H100 +python3 -m sglang.launch_server --model lmsys/gpt-oss-120b-bf16 --tp 8 + +# MXFP4 120B on B200 +python3 -m sglang.launch_server --model openai/gpt-oss-120b --tp 4 + +# BF16 120B on B200 +python3 -m sglang.launch_server --model lmsys/gpt-oss-120b-bf16 --tp 4 +``` + +Benchmark Command + +```bash +python3 -m sglang.bench_one_batch_server --model openai/gpt-oss-120b --base-url http://localhost:30000 --batch-size 32 --input-len 1024 8192 --output-len 512 --show-report +``` + +### Reproduce the evaluation result + +Install gpt-oss + +```bash +git clone https://github.com/openai/gpt-oss.git +cd gpt-oss +pip install -e . +``` + +Evaluation Command + +```bash +DATASET=gpqa +BASE_URL=YOUR_BASE_URL +OPENAI_API_KEY=dummy python -m gpt_oss.evals \ + --base-url ${BASE_URL}/v1 \ + --model dummy \ + --reasoning-effort low,medium,high \ + --eval $DATASET \ + --n-threads 1000 +``` + +### Reproduce the benchmark result of acceptance length +> Note: On B200, if top k is 1, set `--attention-backend trtllm_mha` +```bash +git clone https://github.com/sgl-project/SpecForge.git +cd SpecForge/benchmarks +config_list=( + "1,0,0,0" + "1,3,1,4" + "1,5,4,8" +) +python3 bench_model_speedup.py \ + --model-path openai/gpt-oss-120b \ + --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 \ + --port 20001 \ + --trust-remote-code \ + --mem-fraction-static 0.8 \ + --tp-size 4 \ + --attention-backend fa3 \ + --config-list "${config_list[@]}" \ + --benchmark-list mtbench:80 gsm8k:200 humaneval:200 math500:200 \ + --output lmsys_gpt-oss-120b_Eagle3_result.jsonl + +python3 bench_model_speedup.py \ + --model-path openai/gpt-oss-120b \ + --speculative-draft-model-path nvidia/gpt-oss-120b-Eagle3 \ + --port 20001 \ + --trust-remote-code \ + --mem-fraction-static 0.8 \ + --tp-size 4 \ + --attention-backend fa3 \ + --config-list "${config_list[@]}" \ + --benchmark-list mtbench:80 gsm8k:200 humaneval:200 math500:200 \ + --output nv_gpt-oss-120b_Eagle3_result.jsonl +``` + +### Reproduce the result of speculative decoding speedup + +Launch Command + +```bash +# On Hopper: +# - Tree decoding (topk > 1) and chain decoding (topk = 1) are supported on both FA3 and Triton backends. +python3 -m sglang.launch_server --model openai/gpt-oss-120b --speculative-algorithm EAGLE3 --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --tp 4 +python3 -m sglang.launch_server --model openai/gpt-oss-120b --speculative-algorithm EAGLE3 --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 --speculative-num-steps 5 --speculative-eagle-topk 4 --speculative-num-draft-tokens 8 --tp 4 + +# On Blackwell: +# - Chain decoding (topk = 1) is supported on TRTLLM-MHA backend. Tree decoding (topk > 1) is in progress, stay tuned! +# - Both tree decoding (topk > 1) and chain decoding (topk = 1) are supported on the Triton backend. +python3 -m sglang.launch_server --model openai/gpt-oss-120b --speculative-algo EAGLE3 --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --tp 4 +python3 -m sglang.launch_server --model openai/gpt-oss-120b --speculative-algo EAGLE3 --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 --speculative-num-steps 5 --speculative-eagle-topk 4 --speculative-num-draft-tokens 8 --attention-backend triton --tp 4 +``` + +Benchmark Command + +```bash +config_list=( + "1,0,0,0" + "1,3,1,4" + "1,5,4,8" +) +python3 bench_model_speedup.py \ + --model-path openai/gpt-oss-120b \ + --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 \ + --port 20001 \ + --trust-remote-code \ + --mem-fraction-static 0.8 \ + --tp-size 4 \ + --attention-backend fa3 \ + --config-list "${config_list[@]}" \ + --benchmark-list gsm8k:200 humaneval:200 math500:200 \ + --output lmsys_gpt-oss-120b_Eagle3_result.jsonl +``` + +We can gain the best speedup with the following settings: + +- **1.39x** speedup with the `--speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4` setting. +- **1.52x** speedup with the `--speculative-num-steps 5 --speculative-eagle-topk 4 --speculative-num-draft-tokens 8` setting. diff --git a/benchmark/gsm8k/README.md b/benchmark/gsm8k/README.md new file mode 100644 index 000000000..c110f533c --- /dev/null +++ b/benchmark/gsm8k/README.md @@ -0,0 +1,47 @@ +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 200 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 200 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + + +### Benchmark lmql +``` +CUDA_VISIBLE_DEVICES=0,1 lmql serve-model meta-llama/Llama-2-7b-chat-hf --cuda --port 23000 +``` + +``` +python3 bench_other.py --num-questions 100 --backend lmql --parallel 2 +``` diff --git a/benchmark/gsm8k/bench_other.py b/benchmark/gsm8k/bench_other.py new file mode 100644 index 000000000..6dcb9ad7c --- /dev/null +++ b/benchmark/gsm8k/bench_other.py @@ -0,0 +1,151 @@ +import argparse +import ast +import asyncio +import json +import re +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import download_and_cache_file, dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_one_example(lines, i, include_answer): + ret = "Question: " + lines[i]["question"] + "\nAnswer:" + if include_answer: + ret += " " + lines[i]["answer"] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def main(args): + # Select backend + call_generate = get_call_generate(args) + + # Read data + url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" + filename = download_and_cache_file(url) + lines = list(read_jsonl(filename)) + + # Construct prompts + num_questions = args.num_questions + num_shots = args.num_shots + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + + states = [None] * len(labels) + + # Run requests + if args.backend != "lmql": + # Use thread pool + def get_one_answer(i): + answer = call_generate( + prompt=few_shot_examples + questions[i], + temperature=0, + max_tokens=256, + stop=["Question", "Assistant:", "<|separator|>"], + ) + states[i] = answer + + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + else: + # Use asyncio + async def batched_call(batch_size): + for i in range(0, len(questions), batch_size): + tasks = [] + for q in questions[i : i + batch_size]: + tasks.append( + call_generate( + few_shot_examples + q, + temperature=0, + max_tokens=256, + stop="Question", + ) + ) + rets = await asyncio.gather(*tasks) + for j in range(len(rets)): + states[i + j] = rets[j] + + tic = time.perf_counter() + asyncio.run(batched_call(batch_size=args.parallel)) + latency = time.perf_counter() - tic + + preds = [] + for i in range(len(states)): + preds.append(get_answer_value(states[i])) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + + # Print results + print(f"Accuracy: {acc:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Latency: {latency:.3f} s") + + # Dump results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shots", type=int, default=5) + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/gsm8k/bench_sglang.py b/benchmark/gsm8k/bench_sglang.py new file mode 100644 index 000000000..9cdc2cf84 --- /dev/null +++ b/benchmark/gsm8k/bench_sglang.py @@ -0,0 +1,148 @@ +import argparse +import ast +import json +import os +import re +import time + +import numpy as np + +from sglang.lang.api import set_default_backend +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + dump_bench_raw_result, + select_sglang_backend, +) +from sglang.utils import download_and_cache_file, dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_one_example(lines, i, include_answer): + ret = "Question: " + lines[i]["question"] + "\nAnswer:" + if include_answer: + ret += " " + lines[i]["answer"] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def main(args): + # Select backend + set_default_backend(select_sglang_backend(args)) + + # Read data + data_path = args.data_path + url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" + if not os.path.isfile(data_path): + data_path = download_and_cache_file(url) + lines = list(read_jsonl(data_path)) + + # Construct prompts + num_questions = args.num_questions + num_shots = args.num_shots + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q} for q in questions] + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + @sgl.function + def few_shot_gsm8k(s, question): + s += few_shot_examples + question + s += sgl.gen( + "answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"] + ) + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Run requests + tic = time.perf_counter() + states = few_shot_gsm8k.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + preds = [] + for i in range(len(states)): + preds.append(get_answer_value(states[i]["answer"])) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + + # Compute speed + num_output_tokens = sum( + s.get_meta_info("answer")["completion_tokens"] for s in states + ) + output_throughput = num_output_tokens / latency + + # Print results + print(f"Accuracy: {acc:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Latency: {latency:.3f} s") + print(f"Output throughput: {output_throughput:.3f} token/s") + + # Dump results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + dump_bench_raw_result( + path=args.raw_result_file, + states=states, + preds=preds, + labels=labels, + ) + + with open(args.result_file, "a") as fout: + value = { + "task": "gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shots", type=int, default=5) + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/hellaswag/README.md b/benchmark/hellaswag/README.md new file mode 100644 index 000000000..cb7e65366 --- /dev/null +++ b/benchmark/hellaswag/README.md @@ -0,0 +1,47 @@ +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 200 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend lightllm +``` + + +### Benchmark guidance +``` +CUDA_VISIBLE_DEVICES=0,1 python3 bench_other.py --num-questions 200 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + + +### Benchmark lmql +``` +lmql serve-model meta-llama/Llama-2-7b-chat-hf --cuda --port 23000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend lmql --port 23000 --parallel 1 +``` diff --git a/benchmark/hellaswag/bench_other.py b/benchmark/hellaswag/bench_other.py new file mode 100644 index 000000000..cde0794bb --- /dev/null +++ b/benchmark/hellaswag/bench_other.py @@ -0,0 +1,118 @@ +import argparse +import asyncio +import json +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_select +from sglang.utils import download_and_cache_file, read_jsonl + + +def get_one_example(lines, i, include_answer): + ret = lines[i]["activity_label"] + ": " + lines[i]["ctx"] + " " + if include_answer: + ret += lines[i]["endings"][lines[i]["label"]] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def main(args): + # Select backend + call_select = get_call_select(args) + + # Read data + url = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl" + filename = download_and_cache_file(url) + lines = list(read_jsonl(filename)) + + # Construct prompts + num_questions = args.num_questions + num_shots = args.num_shots + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + choices = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + choices.append(lines[i]["endings"]) + labels.append(lines[i]["label"]) + + preds = [None] * len(labels) + + # Run requests + if args.backend != "lmql": + # Use thread pool + def get_one_answer(i): + preds[i] = call_select( + context=few_shot_examples + questions[i], choices=choices[i] + ) + + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + else: + # Use asyncio + async def batched_call(batch_size): + for i in range(0, len(questions), batch_size): + tasks = [] + for q, c in zip( + questions[i : i + batch_size], choices[i : i + batch_size] + ): + tasks.append(call_select(context=few_shot_examples + q, choices=c)) + rets = await asyncio.gather(*tasks) + for j in range(len(rets)): + preds[i + j] = rets[j] + + tic = time.perf_counter() + asyncio.run(batched_call(batch_size=args.parallel)) + + latency = time.perf_counter() - tic + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + print(f"Latency: {latency:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "hellaswag", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shots", type=int, default=20) + parser.add_argument("--data-path", type=str, default="hellaswag_val.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/hellaswag/bench_sglang.py b/benchmark/hellaswag/bench_sglang.py new file mode 100644 index 000000000..2adce99b8 --- /dev/null +++ b/benchmark/hellaswag/bench_sglang.py @@ -0,0 +1,109 @@ +import argparse +import json +import os +import time + +import numpy as np + +from sglang.lang.api import set_default_backend +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import download_and_cache_file, read_jsonl + + +def get_one_example(lines, i, include_answer): + ret = lines[i]["activity_label"] + ": " + lines[i]["ctx"] + " " + if include_answer: + ret += lines[i]["endings"][lines[i]["label"]] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def main(args): + # Select backend + set_default_backend(select_sglang_backend(args)) + + # Read data + data_path = args.data_path + url = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl" + if not os.path.isfile(data_path): + data_path = download_and_cache_file(url) + lines = list(read_jsonl(data_path)) + + # Construct prompts + num_questions = args.num_questions + num_shots = args.num_shots + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + choices = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + choices.append(lines[i]["endings"]) + labels.append(lines[i]["label"]) + arguments = [{"question": q, "choices": c} for q, c in zip(questions, choices)] + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + @sgl.function + def few_shot_hellaswag(s, question, choices): + s += few_shot_examples + question + s += sgl.select("answer", choices=choices) + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Run requests + tic = time.perf_counter() + rets = few_shot_hellaswag.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + preds = [choices[i].index(rets[i]["answer"]) for i in range(len(rets))] + latency = time.perf_counter() - tic + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + print(f"Latency: {latency:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "hellaswag", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shots", type=int, default=20) + parser.add_argument("--data-path", type=str, default="hellaswag_val.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/hf3fs/bench.sh b/benchmark/hf3fs/bench.sh new file mode 100644 index 000000000..049116b89 --- /dev/null +++ b/benchmark/hf3fs/bench.sh @@ -0,0 +1,59 @@ +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages:/usr/local/lib/python3.12/dist-packages/torch/lib +python3 benchmark/hf3fs/bench_client.py + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages:/usr/local/lib/python3.12/dist-packages/torch/lib +SGLANG_HICACHE_HF3FS_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hf3fs/hf3fs.json \ +python3 benchmark/hf3fs/bench_storage.py + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages:/usr/local/lib/python3.12/dist-packages/torch/lib +export SGLANG_HICACHE_HF3FS_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hf3fs/hf3fs.json +echo '{"file_path_prefix": "/data/hf3fs-test-0", "file_size": 1099511627776, "numjobs": 16, "entries": 8}' > \ +${SGLANG_HICACHE_HF3FS_CONFIG_PATH} +python3 benchmark/hf3fs/bench_zerocopy.py + +#################################################################################################### + +rm -rf nohup.out && \ +nohup python3 -m sglang.launch_server \ + --model-path /code/models/Qwen3-32B/ \ + --host 0.0.0.0 --port 33301 \ + --page-size 64 \ + --enable-hierarchical-cache \ + --hicache-ratio 2 --hicache-size 0 \ + --hicache-write-policy write_through \ + --hicache-storage-backend hf3fs & + +rm -rf bench_multiturn.out && \ +nohup python3 benchmark/hicache/bench_multiturn.py \ + --model-path /code/models/Qwen3-32B \ + --dataset-path /code/models/ShareGPT_V3_unfiltered_cleaned_split.json \ + --port 33301 \ + --request-length 2048 --num-clients 512 --num-rounds 3 --max-parallel 8 \ + > bench_multiturn.out & + +#################################################################################################### + +rm -rf nohup.out && \ +nohup python3 -m sglang.launch_server \ + --model-path /code/models/DeepSeek-R1/ \ + --tp 16 --nnodes 2 --node-rank 0 \ + --dist-init-addr 10.74.249.153:5000 \ + --host 0.0.0.0 --port 33301 \ + --page-size 64 \ + --enable-hierarchical-cache \ + --hicache-ratio 2 --hicache-size 60 \ + --hicache-write-policy write_through \ + --hicache-storage-backend hf3fs & + +rm -rf bench_multiturn.out && \ +nohup python3 benchmark/hicache/bench_multiturn.py \ + --model-path /code/models/Qwen3-32B \ + --dataset-path /code/models/ShareGPT_V3_unfiltered_cleaned_split.json \ + --port 33301 \ + --request-length 2048 --num-clients 1024 --num-rounds 3 --max-parallel 8 \ + > bench_multiturn.out & + +#################################################################################################### + +ps aux | grep "sglang.launch_server" | grep -v grep | awk '{print $2}' | xargs kill -9 +ps aux | grep "bench_multiturn.py" | grep -v grep | awk '{print $2}' | xargs kill -9 diff --git a/benchmark/hf3fs/bench_client.py b/benchmark/hf3fs/bench_client.py new file mode 100644 index 000000000..33c502575 --- /dev/null +++ b/benchmark/hf3fs/bench_client.py @@ -0,0 +1,162 @@ +import concurrent.futures +import logging +import random +import time +from typing import List + +import torch +from tqdm import tqdm + +from sglang.srt.mem_cache.storage.hf3fs.client_hf3fs import Hf3fsClient + + +def print_stats(x: List[int]): + x = sorted(x) + lenx = len(x) + print( + f"mean = {sum(x)/len(x):.2f}, " + f"min = {min(x):.2f}, " + f"p25 = {x[int(lenx*0.25)]:.2f}, " + f"p50 = {x[int(lenx*0.5)]:.2f}, " + f"p75 = {x[int(lenx*0.75)]:.2f}, " + f"max = {max(x):.2f}" + ) + + +def test(): + # /path/to/hf3fs + file_path = "/data/bench.bin" + file_size = 1 << 40 + bytes_per_page = 16 << 20 + entries = 32 + file_ops = Hf3fsClient(file_path, file_size, bytes_per_page, entries) + + print("test batch_read / batch_write") + num_pages = 128 + dtype = torch.bfloat16 + numel = bytes_per_page // dtype.itemsize + offsets = list(range(file_size // bytes_per_page)) + random.shuffle(offsets) + offsets = offsets[:num_pages] + offsets = [i * bytes_per_page for i in offsets] + tensor_writes = [ + torch.randn(numel, dtype=dtype) + for _ in tqdm(range(num_pages), desc="prepare tensor") + ] + for i in tqdm(range(0, num_pages, file_ops.entries), desc="batch_write"): + results = file_ops.batch_write( + offsets[i : i + file_ops.entries], tensor_writes[i : i + file_ops.entries] + ) + assert all([result == numel * dtype.itemsize for result in results]) + tensor_reads = [ + torch.empty(numel, dtype=dtype) + for _ in tqdm(range(num_pages), desc="prepare tensor") + ] + for i in tqdm(range(0, num_pages, file_ops.entries), desc="batch_read"): + results = file_ops.batch_read( + offsets[i : i + file_ops.entries], tensor_reads[i : i + file_ops.entries] + ) + assert all([result == numel * dtype.itemsize for result in results]) + assert all([torch.allclose(r, w) for r, w in zip(tensor_reads, tensor_writes)]) + + file_ops.close() + print("test done") + + +def bench(): + file_path = "/data/bench.bin" + file_size = 1 << 40 + bytes_per_page = 16 << 20 + entries = 8 + numjobs = 16 + + dtype = torch.bfloat16 + numel = bytes_per_page // dtype.itemsize + + file_ops = [ + Hf3fsClient(file_path, file_size, bytes_per_page, entries) + for _ in range(numjobs) + ] + + num_page = entries + + offsets = list(range(file_size // bytes_per_page)) + tensors_write = [torch.randn(numel, dtype=dtype)] * num_page + tensors_read = [torch.empty(numel, dtype=dtype)] * num_page + random.shuffle(offsets) + + warmup = 50 + iteration = 100 + + executor = concurrent.futures.ThreadPoolExecutor(max_workers=numjobs) + + w_bw = [] + w_size = num_page * numjobs * bytes_per_page / (1 << 30) + for i in tqdm(range(warmup + iteration), desc="Benchmarking write (GB/s)"): + _offsets = [ + [ + offset * bytes_per_page + for offset in offsets[ + (i * numjobs + j) * num_page : (i * numjobs + j + 1) * num_page + ] + ] + for j in range(numjobs) + ] + tik = time.perf_counter() + futures = [ + executor.submit(file_ops[j].batch_write, offset, tensors_write) + for j, offset in enumerate(_offsets) + ] + results = [future.result() for future in futures] + tok = time.perf_counter() + if i < warmup: + continue + w_bw.append(w_size / (tok - tik)) + results = [ + _result == bytes_per_page for result in results for _result in result + ] + assert all(results) + print_stats(w_bw) + + r_bw = [] + r_size = w_size + for i in tqdm(range(warmup + iteration), desc="Benchmarking read (GB/s)"): + _offsets = [ + [ + offset * bytes_per_page + for offset in offsets[ + (i * numjobs + j) * num_page : (i * numjobs + j + 1) * num_page + ] + ] + for j in range(numjobs) + ] + tik = time.perf_counter() + futures = [ + executor.submit(file_ops[j].batch_read, offset, tensors_read) + for j, offset in enumerate(_offsets) + ] + results = [future.result() for future in futures] + tok = time.perf_counter() + if i < warmup: + continue + r_bw.append(r_size / (tok - tik)) + results = [ + _result == bytes_per_page for result in results for _result in result + ] + assert all(results) + print_stats(r_bw) + + executor.shutdown(wait=True) + for _file_ops in file_ops: + _file_ops.close() + print("bench done") + + +def main(): + logging.basicConfig(level=logging.INFO) + test() + bench() + + +if __name__ == "__main__": + main() diff --git a/benchmark/hf3fs/bench_storage.py b/benchmark/hf3fs/bench_storage.py new file mode 100644 index 000000000..f0ce171bf --- /dev/null +++ b/benchmark/hf3fs/bench_storage.py @@ -0,0 +1,258 @@ +import json +import logging +import os +import random +import time +from typing import List + +import torch +from tqdm import tqdm + +from sglang.srt.mem_cache.storage.hf3fs.mini_3fs_metadata_server import ( + Hf3fsLocalMetadataClient, +) +from sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs import HiCacheHF3FS + + +def print_stats(x: List[int]): + x = sorted(x) + lenx = len(x) + print( + f"mean = {sum(x)/len(x):.2f}, " + f"min = {min(x):.2f}, " + f"p25 = {x[int(lenx*0.25)]:.2f}, " + f"p50 = {x[int(lenx*0.5)]:.2f}, " + f"p75 = {x[int(lenx*0.75)]:.2f}, " + f"max = {max(x):.2f}" + ) + + +def test(): + # Qwen3-32B + layer_num = 64 + head_num, head_dim = 8, 128 + kv_lora_rank, qk_rope_head_dim = 0, 0 + store_dtype = torch.bfloat16 + tokens_per_page = 64 + + file_path_prefix = "/data/test" + file_size = 128 << 20 + numjobs = 16 + bytes_per_page = 16 << 20 + entries = 2 + dtype = store_dtype + + config_path = os.getenv(HiCacheHF3FS.default_env_var) + assert config_path + try: + with open(config_path, "w") as f: + json.dump( + { + "file_path_prefix": file_path_prefix, + "file_size": file_size, + "numjobs": numjobs, + "entries": entries, + }, + f, + ) + except Exception as e: + raise RuntimeError(f"Failed to dump config to {config_path}: {str(e)}") + hicache_hf3fs = HiCacheHF3FS.from_env_config(bytes_per_page, dtype) + + numel = 2 * tokens_per_page * layer_num * head_num * head_dim + assert numel * dtype.itemsize == bytes_per_page + + num_pages = 10 + tensors = {} + for i in range(num_pages): + k = f"key_{i}" + v = torch.randn((numel,)).to(dtype=dtype) + ok = hicache_hf3fs.set(k, v) + if i < (file_size // bytes_per_page): + assert ok, f"Failed to insert {k}" + else: + assert not ok + tensors[k] = v + assert hicache_hf3fs.get("key_8") is None + assert hicache_hf3fs.get("key_9") is None + + start = 0 + for i in range(start, start + hicache_hf3fs.num_pages): + k = f"key_{i}" + assert hicache_hf3fs.exists(k) + out = hicache_hf3fs.get(k) + assert out is not None + v = tensors[k] + assert torch.allclose(v, out, atol=1e-3), f"Tensor mismatch for {k}" + + assert not hicache_hf3fs.exists("not_exists") + + hicache_hf3fs.delete("key_7") + v2 = torch.randn((numel,)).to(dtype=dtype) + assert hicache_hf3fs.set("key_new", v2) + assert torch.allclose(hicache_hf3fs.get("key_new"), v2, atol=1e-3) + + hicache_hf3fs.clear() + assert ( + len(hicache_hf3fs.metadata_client.rank_metadata.free_pages) + == hicache_hf3fs.metadata_client.rank_metadata.num_pages + ) + + # batch + num_pages = 10 + tensors = {} + keys = [] + values = [] + for i in range(num_pages): + k = f"key_{i}" + keys.append(k) + v = torch.randn((numel,)).to(dtype=dtype) + values.append(v) + + ok = hicache_hf3fs.batch_set(keys, values) + assert not ok + assert hicache_hf3fs.get("key_8") is None + assert hicache_hf3fs.get("key_9") is None + + results = hicache_hf3fs.batch_get(keys[: hicache_hf3fs.num_pages]) + for result, key, value in zip( + results, keys[: hicache_hf3fs.num_pages], values[: hicache_hf3fs.num_pages] + ): + assert torch.allclose(value, result, atol=1e-3), f"Tensor mismatch for {key}" + + hicache_hf3fs.close() + os.remove(hicache_hf3fs.file_path) + + print("All test cases passed.") + + +def bench(): + # Qwen3-32B + layer_num = 64 + head_num, head_dim = 8, 128 + kv_lora_rank, qk_rope_head_dim = 0, 0 + store_dtype = torch.bfloat16 + tokens_per_page = 64 + + file_path = "/data/test.bin" + file_size = 1 << 40 + numjobs = 16 + bytes_per_page = 16 << 20 + entries = 8 + dtype = store_dtype + hicache_hf3fs = HiCacheHF3FS( + rank=0, + file_path=file_path, + file_size=file_size, + numjobs=numjobs, + bytes_per_page=bytes_per_page, + entries=entries, + dtype=dtype, + metadata_client=Hf3fsLocalMetadataClient(), + ) + + numel = 2 * tokens_per_page * layer_num * head_num * head_dim + assert numel * dtype.itemsize == bytes_per_page + + num_page = 128 + values = [torch.randn((numel,)).to(dtype=dtype) for _ in tqdm(range(num_page))] + + warmup = 50 + iteration = 100 + + w_bw = [] + w_size = num_page * bytes_per_page / (1 << 30) + for i in tqdm(range(warmup + iteration), desc="Benchmarking write (GB/s)"): + keys = [f"{j}" for j in range(i * num_page, (i + 1) * num_page)] + tik = time.perf_counter() + ok = hicache_hf3fs.batch_set(keys, values) + tok = time.perf_counter() + if i < warmup: + continue + w_bw.append(w_size / (tok - tik)) + assert ok + print_stats(w_bw) + + r_bw = [] + r_size = num_page * bytes_per_page / (1 << 30) + for i in tqdm(range(warmup + iteration), desc="Benchmarking read (GB/s)"): + keys = random.sample( + list(hicache_hf3fs.metadata_client.rank_metadata.key_to_index.keys()), + num_page, + ) + tik = time.perf_counter() + results = hicache_hf3fs.batch_get(keys) + tok = time.perf_counter() + if i < warmup: + continue + r_bw.append(r_size / (tok - tik)) + assert all([r is not None for r in results]) + print_stats(r_bw) + + hicache_hf3fs.close() + + +def allclose(): + # Qwen3-32B + layer_num = 64 + head_num, head_dim = 8, 128 + kv_lora_rank, qk_rope_head_dim = 0, 0 + store_dtype = torch.bfloat16 + tokens_per_page = 64 + + file_path = "/data/test.bin" + file_size = 1 << 40 + numjobs = 16 + bytes_per_page = 16 << 20 + entries = 8 + dtype = store_dtype + hicache_hf3fs = HiCacheHF3FS( + rank=0, + file_path=file_path, + file_size=file_size, + numjobs=numjobs, + bytes_per_page=bytes_per_page, + entries=entries, + dtype=dtype, + metadata_client=Hf3fsLocalMetadataClient(), + ) + + numel = 2 * tokens_per_page * layer_num * head_num * head_dim + assert numel * dtype.itemsize == bytes_per_page + + num_page = 128 + values = [torch.randn((numel,)).to(dtype=dtype) for _ in tqdm(range(num_page))] + + iteration = 100 + + for i in tqdm(range(iteration), desc="Benchmarking write (GB/s)"): + keys = [f"{j}" for j in range(i * num_page, (i + 1) * num_page)] + ok = hicache_hf3fs.batch_set(keys, values) + assert ok + + read_keys, read_results = [], [] + for i in tqdm(range(iteration), desc="Benchmarking read (GB/s)"): + keys = random.sample( + list(hicache_hf3fs.metadata_client.rank_metadata.key_to_index.keys()), + num_page, + ) + results = hicache_hf3fs.batch_get(keys) + read_keys.extend(keys) + read_results.extend(results) + assert all([r is not None for r in results]) + + for key, result in tqdm(zip(read_keys, read_results)): + assert torch.allclose(values[int(key) % num_page], result, atol=1e-3) + + hicache_hf3fs.close() + + +def main(): + logging.basicConfig(level=logging.INFO) + test() + bench() + allclose() + + +if __name__ == "__main__": + main() diff --git a/benchmark/hf3fs/bench_zerocopy.py b/benchmark/hf3fs/bench_zerocopy.py new file mode 100644 index 000000000..bfa7bff0e --- /dev/null +++ b/benchmark/hf3fs/bench_zerocopy.py @@ -0,0 +1,140 @@ +import threading +import time + +import torch +from tqdm import tqdm + +from sglang.srt.distributed import ( + get_world_group, + init_distributed_environment, + initialize_model_parallel, +) +from sglang.srt.managers.cache_controller import ( + HiCacheController, + PrefetchOperation, + StorageOperation, +) +from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool +from sglang.srt.mem_cache.memory_pool_host import MHATokenToKVPoolHost + +init_distributed_environment( + world_size=1, + rank=0, + distributed_init_method="tcp://127.0.0.1:23456", + local_rank=0, + backend="gloo", +) + +initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, +) + +group = get_world_group().cpu_group + +max_total_num_tokens = 524288 +page_size = 64 +kv_cache_dtype = torch.bfloat16 +layer_num = 64 +head_num, head_dim = 8, 128 +device = "cuda" +hicache_ratio = 2 +hicache_size = 0 +hicache_mem_layout = "page_first" +# hicache_mem_layout = "layer_first" +hicache_write_policy = "write_through" +hicache_io_backend = "kernel" +hicache_storage_backend = "hf3fs" +prefetch_threshold = 256 + +op_size = 1024 +op_num = 16 + +token_to_kv_pool = MHATokenToKVPool( + max_total_num_tokens, + page_size=page_size, + dtype=kv_cache_dtype, + head_num=head_num, + head_dim=head_dim, + layer_num=layer_num, + device=device, + enable_memory_saver=True, +) + +token_to_kv_pool_allocator = TokenToKVPoolAllocator( + max_total_num_tokens, + dtype=kv_cache_dtype, + device=device, + kvcache=token_to_kv_pool, + need_sort=False, +) + +kv_cache = token_to_kv_pool_allocator.get_kvcache() +token_to_kv_pool_host = MHATokenToKVPoolHost( + kv_cache, + hicache_ratio, + hicache_size, + page_size, + hicache_mem_layout, +) + +load_cache_event = threading.Event() +cache_controller = HiCacheController( + token_to_kv_pool_allocator, + token_to_kv_pool_host, + page_size, + group, + load_cache_event=load_cache_event, + write_policy=hicache_write_policy, + io_backend=hicache_io_backend, + storage_backend=hicache_storage_backend, + prefetch_threshold=prefetch_threshold, +) + +operations = [ + StorageOperation( + torch.tensor(list(range(i, i + op_size))), + list(range(i, i + op_size)), + hash_value=[f"{j}" for j in range(i, i + op_size, page_size)], + ) + for i in tqdm(range(0, op_num * op_size, op_size)) +] + +tik = time.monotonic() +if hicache_mem_layout == "page_first": + for operation in operations: + cache_controller.zerocopy_page_backup(operation, batch_size=128) +elif hicache_mem_layout == "layer_first": + for operation in operations: + cache_controller.generic_page_backup(operation, batch_size=128) +tok = time.monotonic() +print(f"{tok-tik:.6f} s") + +operations = [ + PrefetchOperation( + f"{i}", + torch.tensor(list(range(i, i + op_size))), + list(range(i, i + op_size)), + f"{i}", + ) + for i in tqdm(range(0, op_num * op_size, op_size)) +] + +for operation in operations: + operation.hash_value = [ + f"{j}" + for j in range( + int(operation.last_hash), int(operation.last_hash) + op_size, page_size + ) + ] + +tik = time.monotonic() +if hicache_mem_layout == "page_first": + for operation in operations: + cache_controller.zerocopy_page_transfer(operation, batch_size=128) +elif hicache_mem_layout == "layer_first": + for operation in operations: + cache_controller.generic_page_transfer(operation, batch_size=128) +tok = time.monotonic() +print(f"{tok-tik:.6f} s") diff --git a/benchmark/hicache/README.md b/benchmark/hicache/README.md new file mode 100644 index 000000000..f2aa6ff05 --- /dev/null +++ b/benchmark/hicache/README.md @@ -0,0 +1,91 @@ +## Run synthetic multi-turn benchmark + +``` +# SGLang server with radix cache disabled +python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --disable-radix-cache + +# SGLang server with radix cache on and first-come-first-serve policy +python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --schedule-policy fcfs + +# The default SGLang server with radix cache on and long-prefix-match policy +python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 + +# SGLang server with hierarchical radix cache enabled +python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --enable-hierarchical-cache + +``` + +``` +python bench_multiturn.py --model-path Qwen/Qwen2.5-14B-Instruct +``` + +Note: The performance gain of hierarchical caching depends on the ratio of reusable tokens to GPU memory capacity. The more tokens to be reused, the larger the model, and the more constrained the GPU memory size, the greater the benefit one can expect from hierarchical caching. + + +# Benchmark with more datasets +## Download Dataset +```bash +./download.sh {sharegpt|ultragpt|loogle|nextqa|all} +``` +This script will automatically download the required dataset to the current working directory + +## Multiturn Benchmark +### Supported Datasets +- sharegpt +- ultrachat +- loogle +### Example Usage: +```bash +python3 bench_serving.py --model mistralai/Mistral-7B-Instruct-v0.3 --backend sglang \ +--dataset-path longdep_qa.json --dataset-name loogle --request-rate 10 --num-prompts 10 \ +--port 8001 --enable-multiturn --disable-shuffle +``` +This uses `mistralai/Mistral-7B-Instruct-v0.3` model with `sglang` as backend. The dataset +is `longdep_qa.json`. We send `10 conversations` with `10 req/s` to port 8001. We enable +multiturn chat without shuffling the order of conversations (i.e. following the original +order in the dataset file). + +### Note: +The requests of multiple conversations are sent in a round robin fashion. +For example, if we have 3 conversations A, B, C whose rounds are `[2, 3, 4]` correspondingly, +multiturn chat will send the requests to the backend in the following order: `[A1, B1, C1, A2, B2, C2, B3, C3, C4]` +This has implications on the cache reuse patterns: the cache reuse distance is the largest +under this request pattern (which means a prefix-aware local scheduler in the backend can +yield the most benefit compared to a FIFO scheduler) + +## Shared Prefix Benchmark +### Supported Datasets +- loogle +### Example Usage: +```bash +python3 bench_serving.py --model mistralai/Mistral-7B-Instruct-v0.3 --backend sglang \ +--dataset-path longdep_qa.json --dataset-name loogle --request-rate 10 --num-prompts 10 \ +--port 8001 --enable-shared-prefix --disable-shuffle +``` +### Note: +Shared Prefix benchmark sends the questions for the same prompt together. For example, +if we have 3 shared prefix A, B, C, which have [2, 3, 4] questions correspondingly, +the shared prefix benchmark will send the requests to the +backend in the following order: `[A+Q1, A+Q2, B+Q1, B+Q2, B+Q3, C+Q1, C+Q2, C+Q3]`. + + +## Multi Modality Benchmark (WIP) +### Supported Datasets: +- nextqa +### Example Usage: +```bash +Server: +python3 -m sglang.launch_server --model-path lmms-lab/LLaVA-NeXT-Video-7B --tp 2 --dp 1 --port 8001 \ +--host 0.0.0.0 --mem-fraction-static 0.9 --tokenizer-path llava-hf/llava-1.5-7b-hf \ +--json-model-override-args "{\"architectures\": [\"LlavaVidForCausalLM\"], \"model_type\":\"llava\", \"mm_spatial_pool_stride\":2}" + +Client: +python3 bench_serving.py --model lmms-lab/LLaVA-NeXT-Video-7B --backend sglang --dataset-path \ +NExTVideo --dataset-name nextqa --request-rate 10 --num-prompts 1 --disable-shuffle --port 8001 \ --enable-multiturn --max-frames 16 --tokenizer llava-hf/llava-1.5-7b-hf --fixed-output-len 2048 +``` +Note: for the server args, `tokenizer-path`, overriding architecture are necessary. + +## Supported Backend +- sglang (oai) +- vllm (oai) +- lmdeploy (oai) diff --git a/benchmark/hicache/bench_long_context.py b/benchmark/hicache/bench_long_context.py new file mode 100644 index 000000000..eed0ae5dc --- /dev/null +++ b/benchmark/hicache/bench_long_context.py @@ -0,0 +1,101 @@ +import json +import queue +import time + +import requests +from bench_multiturn import ( + ReadyQueue, + WorkloadGenerator, + gen_payload, + log_to_jsonl_file, + parse_args, +) +from tqdm.asyncio import tqdm + +from sglang.bench_serving import get_tokenizer + + +class ContextWorkloadGenerator(WorkloadGenerator): + def __init__(self, args): + # Construct the base URL for requests + self.baseurl = f"http://{args.host}:{args.port}/" + self.url = self.baseurl + "generate" + + self.tokenizer = get_tokenizer(args.model_path) + self.distribution = args.distribution + self.request_rate = args.request_rate + self.start_time = None + self.finished_time = None + + self.sent_requests = 0 + self.completed_requests = 0 + + self.dataset = json.load(open(args.dataset_path)) + num_requests = min(args.num_clients, len(self.dataset["queries"])) + + init_requests = [] + for i in range(num_requests): + context_id = self.dataset["queries"][i]["context"] + init_requests.append( + ( + i, + gen_payload( + self.dataset["contexts"][context_id] + + self.dataset["queries"][i]["question"], + len( + self.tokenizer( + self.dataset["queries"][i]["reference_answer"] + )["input_ids"] + ), + ), + ) + ) + self.ready_queue = ReadyQueue(init_requests=init_requests) + + self.response_queue = queue.Queue() + self.pbar = tqdm(total=num_requests) + self.performance_metrics = { + "ttft": [], + "latency": [], + "itl": [], + "prompt_len": [], + "cached_tokens": [], + "generated_len": [], + } + + self.max_parallel = args.max_parallel + self.logfile = args.log_file + + def response_handler(self): + while True: + try: + client_id, response = self.response_queue.get( + timeout=10 + ) # Block until response is available + if not response.success: + raise ValueError(f"Request failed with error: {response.error}") + self.performance_metrics["ttft"].append(response.ttft) + self.performance_metrics["itl"].extend(response.itl) + self.performance_metrics["latency"].append(response.latency) + self.performance_metrics["prompt_len"].append(response.prompt_len) + self.performance_metrics["cached_tokens"].append(response.cached_tokens) + self.performance_metrics["generated_len"].append(response.generated_len) + self.completed_requests += 1 + + except queue.Empty: + if self.pbar.n == self.pbar.total: + break + + +if __name__ == "__main__": + args = parse_args() + args.num_rounds = 1 + args.max_parallel = 24 + flush_cache_url = f"http://{args.host}:{args.port}/flush_cache" + + for request_rate in [24, 16, 12, 8, 4, 2, 1]: + args.request_rate = request_rate + requests.post(flush_cache_url) + time.sleep(1) + performance_data = ContextWorkloadGenerator(args).run() + log_to_jsonl_file(performance_data, args.log_file, args.tag) diff --git a/benchmark/hicache/bench_mix.py b/benchmark/hicache/bench_mix.py new file mode 100644 index 000000000..cfd25bc40 --- /dev/null +++ b/benchmark/hicache/bench_mix.py @@ -0,0 +1,567 @@ +import argparse +import asyncio +import json +import logging +import os +import queue +import random +import threading +import time +from dataclasses import dataclass +from functools import wraps + +import aiohttp + +from sglang.bench_serving import ( + RequestFuncOutput, + get_tokenizer, + remove_prefix, + sample_random_requests, +) + +# Set up logger +logger = logging.getLogger(__name__) + +# Set up JSONL file for debug logging +debug_log_file = None +# Create a lock for thread-safe debug log writing +debug_log_lock = threading.Lock() + + +def write_debug_log(data): + global debug_log_file + + """Write debug information to a JSONL file""" + if debug_log_file is None: + return + + # Acquire lock for thread-safe writing + with debug_log_lock: + # Write as JSONL (JSON Line format) + debug_log_file.write(json.dumps(data) + "\n") + debug_log_file.flush() + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Script to benchmark concurrent requests to a server." + ) + parser.add_argument( + "--model-path", + type=str, + default="/data/models/Qwen3-0.6B", + help="model path compatible with Hugging Face Transformers", + ) + parser.add_argument( + "--dataset-path", + type=str, + default="/data/models/ShareGPT_V3_unfiltered_cleaned_split/ShareGPT_V3_unfiltered_cleaned_split.json", + help="local dataset to sample tokens from", + ) + parser.add_argument( + "--host", + type=str, + default="localhost", + help="Server hostname or IP (default: localhost)", + ) + parser.add_argument( + "--port", + type=int, + default=30000, + help="Server port (default: 30000)", + ) + parser.add_argument( + "--duration", + type=int, + default=600, + help="Duration to run the benchmark in seconds (default: 300 seconds)", + ) + parser.add_argument( + "--log-level", + type=str, + default="info", + choices=["debug", "info"], + help="Set the logging level (default: info)", + ) + parser.add_argument( + "--debug-log-file", + type=str, + default="debug.log.jsonl", + help="File to write debug logs in JSONL format", + ) + return parser.parse_args() + + +def load_config(): + config_path = os.getenv("CONFIG_PATH") + if not config_path: + raise ValueError("Environment variable 'CONFIG_PATH' is not set.") + + with open(config_path, "r") as f: + config = json.load(f) + + required_keys = [ + "num_rounds", + "num_clients", + "round_ratios", + "mean_new_tokens_per_round", + "mean_return_tokens_per_round", + "mean_inter_round_interval", + ] + + for key in required_keys: + if key not in config: + raise KeyError(f"Missing required configuration key: {key}") + + num_rounds = config["num_rounds"] + assert len(config["round_ratios"]) == num_rounds + assert len(config["mean_new_tokens_per_round"]) == num_rounds + assert len(config["mean_return_tokens_per_round"]) == num_rounds + assert len(config["mean_inter_round_interval"]) == num_rounds + + print(config) + + return config + + +@dataclass +class UserData: + user_id: int + current_round: int + total_rounds: int + prompt: str + return_tokens: int + start: int + + +def synchronized(): + def _decorator(func): + @wraps(func) + def wrapper(self, *args, **kwargs): + with self.lock: + return func(self, *args, **kwargs) + + return wrapper + + return _decorator + + +class UserGenerator: + def __init__(self, config, model_path, dataset_path): + self.tokenizer_path = model_path + self.tokenizer = get_tokenizer(self.tokenizer_path) + self.dataset_path = dataset_path + + self.user_id = 0 + self.lock = threading.Lock() + + self.num_rounds = config["num_rounds"] + + self.cumulative_ratios = [ + sum(config["round_ratios"][: i + 1]) + for i in range(len(config["round_ratios"])) + ] + self.mean_new_tokens_per_round = config["mean_new_tokens_per_round"] + self.mean_return_tokens_per_round = config["mean_return_tokens_per_round"] + self.mean_inter_round_interval = config["mean_inter_round_interval"] + + self.sigma = 100 + self.range_ratio = 0.8 + assert self.range_ratio <= 1 + + self.candidate_inputs = [ + [ + r + for r in sample_random_requests( + input_len=( + self.mean_new_tokens_per_round[i] * (2 - self.range_ratio) + ), + output_len=( + self.mean_return_tokens_per_round[i] * (2 - self.range_ratio) + ), + num_prompts=config["num_clients"], + range_ratio=self.range_ratio / (2 - self.range_ratio), + tokenizer=self.tokenizer, + dataset_path=self.dataset_path, + random_sample=False, + ) + ] + for i in range(self.num_rounds) + ] + + self.multiturn_queue = [] + + self.user_stats = [0 for _ in range(self.num_rounds)] + self.input_stats = [[0, 0] for _ in range(self.num_rounds)] + self.output_stats = [[0, 0] for _ in range(self.num_rounds)] + + def gen(self): + user_id = self.user_id + self.user_id += 1 + + rand_ratio = random.randint(0, self.cumulative_ratios[-1]) + i = len(self.cumulative_ratios) + for idx, cumulative_ratio in enumerate(self.cumulative_ratios): + if rand_ratio >= cumulative_ratio: + continue + else: + i = idx + 1 + break + total_rounds = i + current_round = 0 + + candidate_input = random.sample(self.candidate_inputs[current_round], 1)[0] + self.input_stats[0][0] += candidate_input.prompt_len + self.input_stats[0][1] += 1 + prompt = f"{user_id} " + candidate_input.prompt + return_tokens = int( + random.gauss(self.mean_return_tokens_per_round[current_round], self.sigma) + ) + if return_tokens <= 0: + return_tokens = self.mean_return_tokens_per_round[current_round] + start = 0 + + user_data = UserData( + user_id, current_round, total_rounds, prompt, return_tokens, start + ) + + self.user_stats[total_rounds - 1] += 1 + + return user_data + + @synchronized() + def push(self, user_data, generated_text, len_itl): + self.output_stats[user_data.current_round][0] += len_itl + 1 + self.output_stats[user_data.current_round][1] += 1 + user_data.current_round += 1 + if user_data.current_round >= user_data.total_rounds: + return + + candidate_input = random.sample( + self.candidate_inputs[user_data.current_round], 1 + )[0] + self.input_stats[user_data.current_round][0] += candidate_input.prompt_len + self.input_stats[user_data.current_round][1] += 1 + user_data.prompt += generated_text + candidate_input.prompt + user_data.return_tokens = int( + random.gauss( + self.mean_return_tokens_per_round[user_data.current_round], self.sigma + ) + ) + if user_data.return_tokens <= 0: + user_data.return_tokens = self.mean_return_tokens_per_round[ + user_data.current_round + ] + interval = random.gauss( + self.mean_inter_round_interval[user_data.current_round], self.sigma + ) + if interval <= 0: + interval = self.mean_inter_round_interval[user_data.current_round] + user_data.start = time.perf_counter() + interval + + if len(self.multiturn_queue) == 0: + self.multiturn_queue.append(user_data) + else: + i = len(self.multiturn_queue) + for idx, d in enumerate(self.multiturn_queue): + if user_data.start < d.start: + i = idx + break + self.multiturn_queue.insert(idx, user_data) + + @synchronized() + def pop(self): + if ( + len(self.multiturn_queue) + and time.perf_counter() > self.multiturn_queue[0].start + ): + return self.multiturn_queue.pop(0) + return self.gen() + + +def gen_payload(prompt, output_len): + payload = { + "text": prompt, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": output_len, + "ignore_eos": True, + }, + "stream": True, + "stream_options": {"include_usage": True}, + "lora_path": "", + "return_logprob": False, + "logprob_start_len": -1, + } + return payload + + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) + + +async def async_request_sglang_generate( + user_data, + url, + atomic_counter, +): + """ + Sends a streaming request to the server. Gathers text token-by-token. + """ + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + headers = {} + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + output = RequestFuncOutput() + payload = gen_payload(user_data.prompt, user_data.return_tokens) + write_debug_log({"timestamp": st, "user_data": user_data.__dict__}) + + try: + async with session.post(url=url, json=payload, headers=headers) as response: + if response.status == 200: + prompt_tokens = 0 + cached_tokens = 0 + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + if data.get("text"): + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + prompt_tokens = (data.get("meta_info") or {}).get( + "prompt_tokens", 0 + ) + cached_tokens = (data.get("meta_info") or {}).get( + "cached_tokens", 0 + ) + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text = data["text"] + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.prompt_len = prompt_tokens + output.cached_tokens = cached_tokens + else: + output.error = response.reason or "" + output.success = False + except Exception as e: + output.success = False + output.error = str(e) + print(f"Request failed: {e}") + + atomic_counter.increment(1) + return output + + +class AtomicCounter: + def __init__(self, initial_value=0): + self._value = initial_value + self.lock = threading.Lock() + + @synchronized() + def increment(self, amount=1): + self._value += amount + + @synchronized() + def get(self): + return self._value + + +class WorkloadGenerator: + def __init__(self, args): + config = load_config() + user_generator = UserGenerator( + config, + args.model_path, + args.dataset_path, + ) + + self.url = f"http://{args.host}:{args.port}/generate" + + self.tokenizer = user_generator.tokenizer + self.start_time = None + self.finished_time = None + self.duration = args.duration + self.done = False + + self.sent_requests = 0 + self.completed_requests = 0 + + self.user_generator = user_generator + self.response_queue = queue.Queue() + self.performance_metrics = { + "ttft": [], + "latency": [], + "prompt_len": [], + "cached_tokens": [], + } + self.max_parallel = config["num_clients"] + + self.atomic_counter = AtomicCounter() + + async def handle_request(self, user_data): + try: + response = await async_request_sglang_generate( + user_data, self.url, self.atomic_counter + ) + self.response_queue.put((user_data, response)) + except Exception as e: + print(f"Request failed: {e}") + self.completed_requests += 1 + + def request_sender(self): + async def request_loop(): + while True: + if self.sent_requests - self.completed_requests < self.max_parallel: + new_request = self.user_generator.pop() + if new_request: + asyncio.create_task(self.handle_request(new_request)) + self.sent_requests += 1 + else: + await asyncio.sleep(0.05) + continue + + if time.perf_counter() - self.start_time > self.duration: + self.done = True + break + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(request_loop()) + loop.close() + + def response_handler(self): + while True: + try: + user_data, response = self.response_queue.get(timeout=10) + logger.info( + f"{((time.perf_counter()-self.start_time)/self.duration*100):.2f}%" + ) + if not response.success: + raise ValueError(f"Request failed with error: {response.error}") + + self.user_generator.push( + user_data, response.generated_text, len(response.itl) + ) + self.performance_metrics["ttft"].append(response.ttft) + self.performance_metrics["latency"].append(response.latency) + self.performance_metrics["prompt_len"].append(response.prompt_len) + self.performance_metrics["cached_tokens"].append(response.cached_tokens) + self.completed_requests += 1 + self.finished_time = time.perf_counter() + + except queue.Empty: + if self.done: + break + except ValueError as e: + print(f"Error processing response for client {user_data}: {e}") + continue + + def run(self): + request_thread = threading.Thread(target=self.request_sender, daemon=True) + response_thread = threading.Thread(target=self.response_handler, daemon=True) + + self.start_time = time.perf_counter() + request_thread.start() + response_thread.start() + + request_thread.join() + response_thread.join() + + performance_data = { + "summary": { + "total_requests": len(self.performance_metrics["ttft"]), + "average_ttft": sum(self.performance_metrics["ttft"]) + / len(self.performance_metrics["ttft"]), + "p90_ttft": sorted(self.performance_metrics["ttft"])[ + int(0.9 * len(self.performance_metrics["ttft"])) + ], + "median_ttft": sorted(self.performance_metrics["ttft"])[ + len(self.performance_metrics["ttft"]) // 2 + ], + "average_latency": sum(self.performance_metrics["latency"]) + / len(self.performance_metrics["latency"]), + "p90_latency": sorted(self.performance_metrics["latency"])[ + int(0.9 * len(self.performance_metrics["latency"])) + ], + "median_latency": sorted(self.performance_metrics["latency"])[ + len(self.performance_metrics["latency"]) // 2 + ], + "throughput": self.atomic_counter.get() + / (self.finished_time - self.start_time), + "cache_hit_rate": ( + 0 + if sum(self.performance_metrics["prompt_len"]) == 0 + else sum(self.performance_metrics["cached_tokens"]) + / sum(self.performance_metrics["prompt_len"]) + ), + }, + } + print("All requests completed") + print("Performance metrics summary:") + print(f" Total requests: {performance_data['summary']['total_requests']}") + print(f" Average TTFT: {performance_data['summary']['average_ttft']:.2f}") + print(f" P90 TTFT: {performance_data['summary']['p90_ttft']:.2f}") + print(f" Median TTFT: {performance_data['summary']['median_ttft']:.2f}") + print( + f" Average latency: {performance_data['summary']['average_latency']:.2f}" + ) + print(f" P90 latency: {performance_data['summary']['p90_latency']:.2f}") + print(f" Median latency: {performance_data['summary']['median_latency']:.2f}") + print( + f" Throughput: {performance_data['summary']['throughput']:.2f} requests per second" + ) + print(f" Cache Hit Rate: {performance_data['summary']['cache_hit_rate']:.6f}") + + user_stats = self.user_generator.user_stats + input_stats = self.user_generator.input_stats + output_stats = self.user_generator.output_stats + print(f"round_ratios: {user_stats}") + print( + f"mean_new_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in input_stats]}" + ) + print( + f"mean_return_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in output_stats]}" + ) + return performance_data + + +def main(): + global debug_log_file + + args = parse_args() + if args.log_level == "debug": + logging.basicConfig(level=logging.DEBUG) + logger.info("use log_level debug") + # Initialize debug log file + debug_log_file = open(args.debug_log_file, "w") + else: + logging.basicConfig(level=logging.INFO) + logger.info("use log_level info") + performance_data = WorkloadGenerator(args).run() + + # Close debug log file if it was opened + if debug_log_file: + debug_log_file.close() + + +if __name__ == "__main__": + main() diff --git a/benchmark/hicache/bench_mix.sh b/benchmark/hicache/bench_mix.sh new file mode 100755 index 000000000..5ff6dca94 --- /dev/null +++ b/benchmark/hicache/bench_mix.sh @@ -0,0 +1,42 @@ +#!/bin/bash + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages:/usr/local/lib/python3.12/dist-packages/torch/lib +rm -rf nohup.out && \ +nohup python3 -m sglang.launch_server \ + --attention-backend triton \ + --model-path /code/models/Qwen3-32B/ \ + --log-level info \ + --tp 4 --mem-frac 0.25 \ + --host 0.0.0.0 --port 33301 \ + --enable-metrics --enable-cache-report \ + --page-size 64 \ + --enable-hierarchical-cache \ + --hicache-ratio 2.5 --hicache-size 0 \ + --hicache-io-backend kernel \ + --hicache-mem-layout layer_first \ + --hicache-write-policy write_through \ + & + +################################################## + +export CONFIG_PATH=/tmp/bench_mix_config.json + +# num_clients: Maximum number of concurrent client requests to be simulated +# round_ratios: Distribution of requests across rounds. Given sum(round_ratios) total requests, +# round_ratios[i] denotes the number of requests that will execute for (i+1) rounds +echo '{ + "num_rounds": 10, + "num_clients": 60, + "round_ratios": [50, 25, 15, 15, 10, 10, 9, 8, 7, 6], + "mean_new_tokens_per_round": [1000, 400, 350, 300, 280, 260, 240, 220, 210, 200], + "mean_return_tokens_per_round": [100, 100, 100, 100, 100, 100, 100, 100, 100, 100], + "mean_inter_round_interval": [30, 30, 30, 30, 30, 30, 30, 30, 30, 30] +}' > ${CONFIG_PATH} + +rm -rf bench_mix.out && \ +nohup python3 /sgl-workspace/sglang/benchmark/hicache/bench_mix.py \ + --model-path /code/models/Qwen3-32B/ \ + --dataset-path /code/models/ShareGPT_V3_unfiltered_cleaned_split.json \ + --port 33301 \ + --duration 600 \ +> bench_mix.out & diff --git a/benchmark/hicache/bench_multiturn.py b/benchmark/hicache/bench_multiturn.py new file mode 100644 index 000000000..a3e8b0d74 --- /dev/null +++ b/benchmark/hicache/bench_multiturn.py @@ -0,0 +1,517 @@ +import argparse +import asyncio +import json +import queue +import random +import threading +import time +from datetime import datetime +from typing import Optional + +import aiohttp +import numpy as np +import requests +from tqdm.asyncio import tqdm + +from sglang.bench_serving import ( + RequestFuncOutput, + get_tokenizer, + remove_prefix, + sample_random_requests, +) + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Script to benchmark concurrent requests to a server." + ) + parser.add_argument( + "--num-clients", + type=int, + default=256, + help="Number of concurrent clients", + ) + parser.add_argument( + "--max-parallel", + type=int, + default=128, + help="Maximum number of parallel requests", + ) + parser.add_argument( + "--request-length", + type=int, + default=512, + help="Length of each new request", + ) + parser.add_argument( + "--output-length", + type=int, + default=64, + help="Length of each output", + ) + parser.add_argument( + "--num-rounds", + type=int, + default=5, + help="Number of rounds per client", + ) + parser.add_argument( + "--distribution", + type=str, + default="poisson", + choices=["poisson", "uniform"], + help="Distribution type for request intervals (poisson or uniform)", + ) + parser.add_argument( + "--request-rate", + type=float, + default=1.0, + help="Average number of requests per second", + ) + parser.add_argument( + "--host", + type=str, + default="localhost", + help="Server hostname or IP (default: localhost)", + ) + parser.add_argument( + "--port", + type=int, + default=30000, + help="Server port (default: 30000)", + ) + parser.add_argument( + "--model-path", + type=str, + default="meta-llama/Llama-3.1-8B-Instruct", + help="model path compatible with Hugging Face Transformers", + ) + parser.add_argument( + "--dataset-path", + type=str, + default="", + help="local dataset to sample tokens from", + ) + parser.add_argument( + "--log-file", + type=str, + default="performance_metrics.jsonl", + help="File to log performance metrics", + ) + parser.add_argument( + "--disable-auto-run", + action="store_true", + help="If set, disable automatically testing with a range of request rates.", + ) + + parser.add_argument( + "--disable-random-sample", + action="store_true", + help="If set, disable random sampling of requests from the ShareGPT dataset.", + ) + parser.add_argument( + "--sub-question-input-length", + type=int, + default=0, + help="Length of the sub question input for each request, if set 0 use request_length", + ) + parser.add_argument( + "--ready-queue-policy", + type=str, + default="random", + help="Policy for popping requests from the ready queue (random or fifo)", + ) + parser.add_argument( + "--tag", + type=str, + default="", + help="Tag of a certain run in the log file", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--lora-path", + type=str, + default="", + help="String of LoRA path. Currently we only support benchmarking on a single LoRA adaptor.", + ) + return parser.parse_args() + + +async def async_request_sglang_generate( + payload, + url, + pbar: Optional[tqdm] = None, +): + """ + Sends a streaming request to the server. Gathers text token-by-token. + """ + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + headers = {} + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + output = RequestFuncOutput() + + try: + async with session.post(url=url, json=payload, headers=headers) as response: + if response.status == 200: + prompt_tokens = 0 + cached_tokens = 0 + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + if data["text"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + prompt_tokens = (data.get("meta_info") or {}).get( + "prompt_tokens", 0 + ) + cached_tokens = (data.get("meta_info") or {}).get( + "cached_tokens", 0 + ) + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text = data["text"] + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.prompt_len = prompt_tokens + output.cached_tokens = cached_tokens + output.generated_len = len(output.itl) + 1 + else: + output.error = response.reason or "" + output.success = False + except Exception as e: + output.success = False + output.error = str(e) + print(f"Request failed: {e}") + + if pbar: + pbar.update(1) + return output + + +def gen_payload(prompt, output_len, lora_path=""): + payload = { + "text": prompt, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": output_len, + "ignore_eos": True, + }, + "stream": True, + "stream_options": {"include_usage": True}, + "lora_path": lora_path, + "return_logprob": False, + "logprob_start_len": -1, + } + return payload + + +def log_to_jsonl_file(data, file_path="performance_metrics.jsonl", tag=""): + """Append the data with a timestamp and tag to the specified JSONL file.""" + timestamped_data = {"timestamp": datetime.now().isoformat(), "tag": tag, **data} + try: + with open(file_path, "a") as file: + file.write( + json.dumps(timestamped_data) + "\n" + ) # Write as a single line in JSONL format + except IOError as e: + print(f"Error writing to JSONL file: {e}") + + +class ReadyQueue: + """ + Thread-safe queue that can pop requests in different orders based on given policy. + """ + + def __init__(self, init_requests=None, policy="random"): + self.lock = threading.Lock() + self.requests = init_requests or [] + self.policy = policy + + def append(self, item): + with self.lock: + self.requests.append(item) + + def pop(self): + with self.lock: + if not self.requests: + return None + if self.policy == "random": + index = random.randrange(len(self.requests)) + return self.requests.pop(index) + elif self.policy == "fifo": + return self.requests.pop(0) + else: + # todo, varying thinking time of clients + raise ValueError(f"{self.policy} not implemented") + + +class WorkloadGenerator: + def __init__(self, args): + # Construct the base URL for requests + self.url = f"http://{args.host}:{args.port}/generate" + + self.tokenizer = get_tokenizer(args.model_path) + self.distribution = args.distribution + self.request_rate = args.request_rate + self.start_time = None + self.finished_time = None + + self.sent_requests = 0 + self.completed_requests = 0 + + self.candidate_inputs = sample_random_requests( + input_len=args.request_length, + output_len=args.output_length, + num_prompts=args.num_clients, + range_ratio=1.0, + tokenizer=self.tokenizer, + dataset_path=args.dataset_path, + random_sample=not args.disable_random_sample, + ) + self.candidate_inputs = [i.prompt for i in self.candidate_inputs] + + if args.sub_question_input_length != 0: + sub_question_input_length = args.sub_question_input_length + else: + sub_question_input_length = args.request_length + + self.sub_question_inputs = sample_random_requests( + input_len=sub_question_input_length, + output_len=args.output_length, + num_prompts=args.num_clients * max(args.num_rounds - 1, 1), + range_ratio=1.0, + tokenizer=self.tokenizer, + dataset_path=args.dataset_path, + random_sample=not args.disable_random_sample, + ) + + init_requests = [ + ( + i, + gen_payload( + self.candidate_inputs[i], args.output_length, args.lora_path + ), + ) + for i in range(args.num_clients) + ] + self.client_records = { + i: {"round": 0, "history": init_requests[i][1]["text"]} + for i in range(args.num_clients) + } + self.ready_queue = ReadyQueue( + init_requests=init_requests, policy=args.ready_queue_policy + ) + self.candidate_inputs = self.candidate_inputs[args.num_clients :] + + self.response_queue = queue.Queue() + self.pbar = tqdm(total=args.num_clients * args.num_rounds) + self.performance_metrics = { + "ttft": [], + "latency": [], + "prompt_len": [], + "cached_tokens": [], + "generated_len": [], + } + self.num_rounds = args.num_rounds + self.max_parallel = args.max_parallel + self.output_length = args.output_length + + async def handle_request(self, item): + try: + client_id, payload = item + response = await async_request_sglang_generate(payload, self.url, self.pbar) + if self.pbar.n == self.pbar.total: + self.finished_time = time.perf_counter() + self.response_queue.put((client_id, response)) + except Exception as e: + print(f"Request failed: {e}") + + def request_sender(self): + async def request_loop(): + while True: + if self.sent_requests - self.completed_requests < self.max_parallel: + new_request = self.ready_queue.pop() + if new_request: + asyncio.create_task(self.handle_request(new_request)) + self.sent_requests += 1 + else: + await asyncio.sleep(0.05) + continue + + if self.pbar.n == self.pbar.total: + break + + # Calculate Poisson-distributed wait time + if self.distribution == "poisson": + sleep_time = random.expovariate(self.request_rate) + elif self.distribution == "uniform": + avg_interval = ( + 1.0 / self.request_rate if self.request_rate > 0 else 1.0 + ) + sleep_time = random.uniform(0, 2 * avg_interval) + else: + raise ValueError("Invalid distribution type") + await asyncio.sleep(sleep_time) # Wait before sending the next request + + # Create and run the event loop for asynchronous requests + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(request_loop()) + loop.close() + + def response_handler(self): + while True: + try: + client_id, response = self.response_queue.get( + timeout=10 + ) # Block until response is available + if not response.success: + raise ValueError(f"Request failed with error: {response.error}") + self.client_records[client_id]["history"] += response.generated_text + self.client_records[client_id]["round"] += 1 + self.performance_metrics["ttft"].append(response.ttft) + self.performance_metrics["latency"].append(response.latency) + self.performance_metrics["prompt_len"].append(response.prompt_len) + self.performance_metrics["cached_tokens"].append(response.cached_tokens) + self.performance_metrics["generated_len"].append(response.generated_len) + self.completed_requests += 1 + + if self.client_records[client_id]["round"] < self.num_rounds: + # append new request to client's history + self.client_records[client_id][ + "history" + ] += self.sub_question_inputs.pop().prompt + self.ready_queue.append( + ( + client_id, + gen_payload( + self.client_records[client_id]["history"], + self.output_length, + args.lora_path, + ), + ) + ) + except queue.Empty: + if self.pbar.n == self.pbar.total: + break + except ValueError as e: + print(f"Error processing response for client {client_id}: {e}") + continue + + def run(self): + request_thread = threading.Thread(target=self.request_sender, daemon=True) + response_thread = threading.Thread(target=self.response_handler, daemon=True) + + self.start_time = time.perf_counter() + request_thread.start() + response_thread.start() + + request_thread.join() + response_thread.join() + self.pbar.close() + + duration = self.finished_time - self.start_time + performance_data = { + "summary": { + "total_requests": len(self.performance_metrics["ttft"]), + "request_rate": self.request_rate, + "average_ttft": sum(self.performance_metrics["ttft"]) + / len(self.performance_metrics["ttft"]), + "p90_ttft": sorted(self.performance_metrics["ttft"])[ + int(0.9 * len(self.performance_metrics["ttft"])) + ], + "median_ttft": sorted(self.performance_metrics["ttft"])[ + len(self.performance_metrics["ttft"]) // 2 + ], + "average_latency": sum(self.performance_metrics["latency"]) + / len(self.performance_metrics["latency"]), + "p90_latency": sorted(self.performance_metrics["latency"])[ + int(0.9 * len(self.performance_metrics["latency"])) + ], + "median_latency": sorted(self.performance_metrics["latency"])[ + len(self.performance_metrics["latency"]) // 2 + ], + "input_token_throughput": sum(self.performance_metrics["prompt_len"]) + / duration, + "output_token_throughput": sum( + self.performance_metrics["generated_len"] + ) + / duration, + "throughput": self.pbar.total / duration, + "cache_hit_rate": ( + 0 + if sum(self.performance_metrics["prompt_len"]) == 0 + else sum(self.performance_metrics["cached_tokens"]) + / sum(self.performance_metrics["prompt_len"]) + ), + }, + } + print("All requests completed") + print("Performance metrics summary:") + print( + f" Total requests: {performance_data['summary']['total_requests']} at {performance_data['summary']['request_rate']} requests per second" + ) + print(f" Average TTFT: {performance_data['summary']['average_ttft']:.2f}") + print(f" P90 TTFT: {performance_data['summary']['p90_ttft']:.2f}") + print(f" Median TTFT: {performance_data['summary']['median_ttft']:.2f}") + print( + f" Average latency: {performance_data['summary']['average_latency']:.2f}" + ) + print(f" P90 latency: {performance_data['summary']['p90_latency']:.2f}") + print(f" Median latency: {performance_data['summary']['median_latency']:.2f}") + print( + f" Input token throughput: {performance_data['summary']['input_token_throughput']:.2f} tokens per second" + ) + print( + f" Output token throughput: {performance_data['summary']['output_token_throughput']:.2f} tokens per second" + ) + print( + f" Request Throughput: {performance_data['summary']['throughput']:.2f} requests per second" + ) + print(f" Cache Hit Rate: {performance_data['summary']['cache_hit_rate']:.6f}") + return performance_data + + +if __name__ == "__main__": + args = parse_args() + flush_cache_url = f"http://{args.host}:{args.port}/flush_cache" + + random.seed(args.seed) + np.random.seed(args.seed) + + if args.disable_auto_run: + print("Running with specified request rate...") + request_rates = [args.request_rate] + else: + print("Auto-running with different request rates...") + request_rates = [16, 14, 12, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1] + + for rate in request_rates: + args.request_rate = rate + requests.post(flush_cache_url) + time.sleep(1) + performance_data = WorkloadGenerator(args).run() + log_to_jsonl_file(performance_data, args.log_file, tag=args.tag) diff --git a/benchmark/hicache/bench_serving.py b/benchmark/hicache/bench_serving.py new file mode 100644 index 000000000..e38d0d0ea --- /dev/null +++ b/benchmark/hicache/bench_serving.py @@ -0,0 +1,1029 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/backend_request_func.py +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/benchmark_serving.py + +""" +Benchmark online serving with dynamic requests. + +Usage: +python3 -m sglang.bench_serving --backend sglang --num-prompt 10 + +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --request-rate-range 1,2,4,8,16,32 --random-input 4096 --random-output 1024 --random-range-ratio 0.125 --multi +""" + +import argparse +import asyncio +import json +import os +import random +import sys +import time +import traceback +import warnings +from argparse import ArgumentParser +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple + +import aiohttp +import numpy as np +import requests +from data_processing import MsgContent, SampleOutput, get_dataset +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +from sglang.bench_serving import get_tokenizer, remove_prefix, set_ulimit + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) + +global args + + +@dataclass +class RequestFuncInput: + prompts: List[Tuple[MsgContent, int, int]] + api_url: str + model: str + lora_name: str + extra_request_body: Dict[str, Any] + + # For multiturn chat, store the context + prev_messages: List = field(default_factory=list) + finished_prompts: int = 0 + + +@dataclass +class RequestFuncOutput: + generated_text: List[str] = field(default_factory=list) + prompt_len: List[int] = field(default_factory=list) + output_len: List[int] = field(default_factory=list) + latency: List[float] = field(default_factory=list) + ttft: List[float] = field(default_factory=list) + itl: List[float] = field(default_factory=list) # List of inter-token latencies + + success: bool = False + error: str = "" + + +# set ignore_eos True by default +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + queue: asyncio.Queue, + tokenizer: PreTrainedTokenizerBase, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + "completions" + ), "OpenAI Completions API URL must end with 'completions'." + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "model": request_func_input.model, + "temperature": 0.0, + "best_of": 1, + "stream": not args.disable_stream, + "stream_options": {"include_usage": True}, + "ignore_eos": not args.disable_ignore_eos, + **request_func_input.extra_request_body, + } + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + } + + output = RequestFuncOutput() + + prompt_idx = request_func_input.finished_prompts + messages = request_func_input.prev_messages + prompt, input_len, max_tokens = request_func_input.prompts[prompt_idx] + prompt_len = sum( + prompt[1] + prompt[2] # input_len + output_len + for prompt in request_func_input.prompts[:prompt_idx] + ) + prompt_len += input_len + + # Messages + messages.append( + { + "role": "user", + "content": prompt, + } + ) + payload["messages"] = messages + payload["max_tokens"] = max_tokens + + # output.prompt_len = request_func_input.prompt_len + # print(payload) + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + actual_prompt_len = prompt_len - 1 + actual_output_len = 0 + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + timestamp = time.perf_counter() + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["usage"] is not None and len(data["usage"]) > 0: + actual_prompt_len = data["usage"]["prompt_tokens"] + actual_output_len = data["usage"]["completion_tokens"] + continue + delta = data["choices"][0]["delta"] + + if delta.get("content", None): + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft.append(ttft) + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + generated_text += delta["content"] + most_recent_timestamp = timestamp + + output.prompt_len.append(actual_prompt_len) # truncate + output.output_len.append(actual_output_len) + output.generated_text.append(generated_text) + output.success = True + output.latency.append(latency) + + # Prepare for the new request + request_func_input.prompts[prompt_idx] = ( + prompt, + input_len, + actual_output_len, # changes from max_tokens to output_len + ) + prompt_idx += 1 + messages.append( + { + "role": "assistant", + "content": generated_text, + } + ) + + # Move the new request to the end of the queue + if prompt_idx < len(request_func_input.prompts): + request_func_input.finished_prompts = prompt_idx + request_func_input.prev_messages = messages + await queue.put(request_func_input) + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_profile(api_url: str) -> RequestFuncOutput: + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + output = RequestFuncOutput() + try: + async with session.post(url=api_url) as response: + if response.status == 200: + output.success = True + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + return output + + +ASYNC_REQUEST_FUNCS = { + "sglang": async_request_openai_completions, + "vllm": async_request_openai_completions, + "lmdeploy": async_request_openai_completions, +} + + +@dataclass +class BenchmarkMetrics: + completed: int + total_input: int + total_output: int + total_output_retokenized: int + request_throughput: float + input_throughput: float + output_throughput: float + output_throughput_retokenized: float + total_throughput: float + total_throughput_retokenized: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + p90_ttft_ms: float + p99_ttft_ms: float + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + p90_tpot_ms: float + p99_tpot_ms: float + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + p90_itl_ms: float + p99_itl_ms: float + mean_e2e_latency_ms: float + median_e2e_latency_ms: float + std_e2e_latency_ms: float + p99_e2e_latency_ms: float + concurrency: float + + +async def get_requests( + input_requests_queue: asyncio.Queue, + request_rate: float, + num_actual_requests: int, +) -> AsyncGenerator[RequestFuncInput, None]: + for _ in range(num_actual_requests): + try: + request = await asyncio.wait_for( + input_requests_queue.get(), timeout=300 + ) # Wait for 5 minutes then abort + except Exception as e: + print(f"exception: {e}") + break + + yield request + + if request_rate == float("inf"): + continue + + interval = np.random.exponential(1.0 / request_rate) + await asyncio.sleep(interval) + + +def calculate_metrics( + outputs: List[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase, + backend: str, +) -> Tuple[BenchmarkMetrics, List[int]]: + output_lens: List[int] = [] + retokenized_output_lens: List[int] = [] + total_input = 0 + completed = 0 + itls: List[float] = [] + tpots: List[float] = [] + ttfts: List[float] = [] + e2e_latencies: List[float] = [] + output_success = 0 + for i in range(len(outputs)): + if outputs[i].success: + output_success += 1 + assert len(outputs[i].generated_text) == len(outputs[i].latency) + assert len(outputs[i].generated_text) == len(outputs[i].ttft) + for j in range(len(outputs[i].generated_text)): + output_len = outputs[i].output_len[j] + output_lens.append(output_len) + retokenized_output_len = len( + tokenizer.encode( + outputs[i].generated_text[j], add_special_tokens=False + ) + ) + retokenized_output_lens.append(retokenized_output_len) + total_input += outputs[i].prompt_len[j] + if output_len > 1: + tpots.append( + (outputs[i].latency[j] - outputs[i].ttft[j]) / (output_len - 1) + ) + + completed += 1 + itls += outputs[i].itl + ttfts += outputs[i].ttft + e2e_latencies += outputs[i].latency + + else: + output_lens.append(0) + retokenized_output_lens.append(0) + + if completed == 0: + warnings.warn( + "All requests failed. This is likely due to a misconfiguration " + "on the benchmark arguments.", + stacklevel=2, + ) + metrics = BenchmarkMetrics( + completed=completed, + total_input=total_input, + total_output=sum(output_lens), + total_output_retokenized=sum(retokenized_output_lens), + request_throughput=completed / dur_s, + input_throughput=total_input / dur_s, + output_throughput=sum(output_lens) / dur_s, + output_throughput_retokenized=sum(retokenized_output_lens) / dur_s, + total_throughput=(total_input + sum(output_lens)) / dur_s, + total_throughput_retokenized=(total_input + sum(retokenized_output_lens)) + / dur_s, + mean_ttft_ms=np.mean(ttfts or 0) + * 1000, # ttfts is empty if streaming is not supported by backend + median_ttft_ms=np.median(ttfts or 0) * 1000, + std_ttft_ms=np.std(ttfts or 0) * 1000, + p90_ttft_ms=np.percentile(ttfts or 0, 90) * 1000, + p99_ttft_ms=np.percentile(ttfts or 0, 99) * 1000, + mean_tpot_ms=np.mean(tpots or 0) * 1000, + median_tpot_ms=np.median(tpots or 0) * 1000, + std_tpot_ms=np.std(tpots or 0) * 1000, + p90_tpot_ms=np.percentile(tpots or 0, 90) * 1000, + p99_tpot_ms=np.percentile(tpots or 0, 99) * 1000, + mean_itl_ms=np.mean(itls or 0) * 1000, + median_itl_ms=np.median(itls or 0) * 1000, + std_itl_ms=np.std(itls or 0) * 1000, + p90_itl_ms=np.percentile(itls or 0, 90) * 1000, + p99_itl_ms=np.percentile(itls or 0, 99) * 1000, + mean_e2e_latency_ms=np.mean(e2e_latencies) * 1000, + median_e2e_latency_ms=np.median(e2e_latencies) * 1000, + std_e2e_latency_ms=np.std(e2e_latencies) * 1000, + p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000, + concurrency=np.sum(e2e_latencies) / dur_s, + ) + + return metrics, output_lens + + +async def benchmark( + backend: str, + api_url: str, + base_url: str, + model_id: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: SampleOutput, + request_rate: float, + max_concurrency: Optional[int], + disable_tqdm: bool, + lora_name: str, + extra_request_body: Dict[str, Any], + profile: bool, + enable_shared_prefix: bool, +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + # Limit concurrency + # From https://github.com/vllm-project/vllm/pull/9390 + semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None + + async def limited_request_func(request_func_input, queue, tokenizer, pbar): + if semaphore is None: + return await request_func( + request_func_input=request_func_input, + queue=queue, + tokenizer=tokenizer, + pbar=pbar, + ) + async with semaphore: + return await request_func( + request_func_input=request_func_input, + queue=queue, + tokenizer=tokenizer, + pbar=pbar, + ) + + num_actual_requests = sum(len(r) for r in input_requests) + print(f"Num of shared prefixes or conversations: {len(input_requests)}") + print(f"Num of total requests: {num_actual_requests}") + + # flatten the requests for shared prefix + if enable_shared_prefix: + input_requests = [[r] for requests in input_requests for r in requests] + inputs_requests_queue = asyncio.Queue(maxsize=len(input_requests)) + print("Starting initial single prompt test run...") + # NOTE: Just use the first request of the first conversation for warmup + test_input = RequestFuncInput( + model=model_id, + prompts=input_requests[0][:1], + api_url=api_url, + lora_name=lora_name, + extra_request_body=extra_request_body, + ) + test_output = await request_func( + request_func_input=test_input, queue=inputs_requests_queue, tokenizer=tokenizer + ) + if not test_output.success: + raise ValueError( + "Initial test run failed - Please make sure benchmark arguments " + f"are correctly specified. Error: {test_output.error}" + ) + else: + print("Initial test run completed. Starting main benchmark run...") + + # Check the states + assert inputs_requests_queue.empty() + + # Flush cache + if "sglang" in backend: + requests.post(base_url + "/flush_cache") + + time.sleep(1.0) + + # Start profiler + if profile: + print("Starting profiler...") + profile_output = await async_request_profile( + api_url=base_url + "/start_profile" + ) + if profile_output.success: + print("Profiler started") + + for request in input_requests: + request_func_input = RequestFuncInput( + model=model_id, + prompts=request, + api_url=api_url, + lora_name=lora_name, + extra_request_body=extra_request_body, + ) + inputs_requests_queue.put_nowait(request_func_input) + if ( + not args.enable_multiturn + and not args.enable_shared_prefix + and not args.dataset_name == "generated-shared-prefix" + ): + assert len(input_requests) == num_actual_requests + + pbar = None if disable_tqdm else tqdm(total=num_actual_requests) + + benchmark_start_time = time.perf_counter() + tasks: List[asyncio.Task] = [] + async for request in get_requests( + inputs_requests_queue, request_rate, num_actual_requests + ): + tasks.append( + asyncio.create_task( + limited_request_func( + request_func_input=request, + queue=inputs_requests_queue, + tokenizer=tokenizer, + pbar=pbar, + ) + ) + ) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + # Stop profiler + if profile: + print("Stopping profiler...") + profile_output = await async_request_profile(api_url=base_url + "/stop_profile") + if profile_output.success: + print("Profiler stopped") + + if pbar is not None: + pbar.close() + + # Compute metrics and print results + benchmark_duration = time.perf_counter() - benchmark_start_time + metrics, output_lens = calculate_metrics( + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + backend=backend, + ) + + print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) + print("{:<40} {:<10}".format("Backend:", backend)) + print("{:<40} {:<10}".format("Traffic request rate:", request_rate)) + print( + "{:<40} {:<10}".format( + "Max request concurrency:", + max_concurrency if max_concurrency else "not set", + ) + ) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) + print( + "{:<40} {:<10}".format( + "Total generated tokens (retokenized):", metrics.total_output_retokenized + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", metrics.request_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", metrics.input_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", metrics.output_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Total token throughput (tok/s):", metrics.total_throughput + ) + ) + print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency)) + print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) + print( + "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) + ) + print( + "{:<40} {:<10.2f}".format( + "Median E2E Latency (ms):", metrics.median_e2e_latency_ms + ) + ) + print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) + print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) + print("{:<40} {:<10.2f}".format("P90 TTFT (ms):", metrics.p90_ttft_ms)) + print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) + print( + "{s:{c}^{n}}".format(s="Time per Output Token (excl. 1st token)", n=50, c="-") + ) + print("{:<40} {:<10.2f}".format("Mean TPOT (ms):", metrics.mean_tpot_ms)) + print("{:<40} {:<10.2f}".format("Median TPOT (ms):", metrics.median_tpot_ms)) + print("{:<40} {:<10.2f}".format("P90 TPOT (ms):", metrics.p90_tpot_ms)) + print("{:<40} {:<10.2f}".format("P99 TPOT (ms):", metrics.p99_tpot_ms)) + print("{s:{c}^{n}}".format(s="Inter-token Latency", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) + print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) + print("{:<40} {:<10.2f}".format("P90 ITL (ms):", metrics.p90_itl_ms)) + print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) + print("=" * 50) + + if ( + metrics.median_ttft_ms is not None + and metrics.mean_itl_ms is not None + and metrics.output_throughput is not None + ): + result = { + # Arguments + "backend": args.backend, + "dataset_name": args.dataset_name, + "request_rate": request_rate, + "max_concurrency": max_concurrency, + "fixed_output_len": args.fixed_output_len, + "random_input_len": args.random_input_len, + "random_output_len": args.random_output_len, + "random_range_ratio": args.random_range_ratio, + # Results + "duration": benchmark_duration, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + "std_e2e_latency_ms": metrics.std_e2e_latency_ms, + "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms, + "mean_ttft_ms": metrics.mean_ttft_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "std_ttft_ms": metrics.std_ttft_ms, + "p99_ttft_ms": metrics.p99_ttft_ms, + "mean_tpot_ms": metrics.mean_tpot_ms, + "median_tpot_ms": metrics.median_tpot_ms, + "std_tpot_ms": metrics.std_tpot_ms, + "p99_tpot_ms": metrics.p99_tpot_ms, + "mean_itl_ms": metrics.mean_itl_ms, + "median_itl_ms": metrics.median_itl_ms, + "std_itl_ms": metrics.std_itl_ms, + "p99_itl_ms": metrics.p99_itl_ms, + "concurrency": metrics.concurrency, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "fixed_output_len": args.fixed_output_len, + "random_input_len": args.random_input_len, + "random_output_len": args.random_output_len, + "random_range_ratio": args.random_range_ratio, + "duration": benchmark_duration, + "completed": metrics.completed, + } + else: + print(f"Error running benchmark for request rate: {request_rate}") + print("-" * 30) + + # Determine output file name + if args.output_file: + output_file_name = args.output_file + else: + now = datetime.now().strftime("%m%d") + if args.dataset_name == "random": + output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" + else: + output_file_name = ( + f"{args.backend}_{now}_{args.num_prompts}_{args.dataset_name}.jsonl" + ) + + # Append results to a JSONL file + with open(output_file_name, "a") as file: + file.write(json.dumps(result) + "\n") + + result = { + "duration": benchmark_duration, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "mean_ttft_ms": metrics.mean_ttft_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "std_ttft_ms": metrics.std_ttft_ms, + "p90_ttft_ms": metrics.p90_ttft_ms, + "p99_ttft_ms": metrics.p99_ttft_ms, + "mean_tpot_ms": metrics.mean_tpot_ms, + "median_tpot_ms": metrics.median_tpot_ms, + "std_tpot_ms": metrics.std_tpot_ms, + "p90_tpot_ms": metrics.p90_tpot_ms, + "p99_tpot_ms": metrics.p99_tpot_ms, + "mean_itl_ms": metrics.mean_itl_ms, + "median_itl_ms": metrics.median_itl_ms, + "std_itl_ms": metrics.std_itl_ms, + "p90_itl_ms": metrics.p90_itl_ms, + "p99_itl_ms": metrics.p99_itl_ms, + "input_lens": [output.prompt_len for output in outputs], + "output_lens": output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + } + return result + + +def run_benchmark(args_: argparse.Namespace): + global args + args = args_ + + # Set default value for max_concurrency if not present + if not hasattr(args, "max_concurrency"): + args.max_concurrency = None + + # Set global environments + set_ulimit() + random.seed(args.seed) + np.random.seed(args.seed) + + extra_request_body = {} + if args.extra_request_body: + extra_request_body = json.loads(args.extra_request_body) + + # Set url + if args.port is None: + args.port = { + "sglang": 30000, + "lmdeploy": 23333, + "vllm": 8000, + }.get(args.backend, 30000) + + model_url = ( + f"{args.base_url}/v1/models" + if args.base_url + else f"http://{args.host}:{args.port}/v1/models" + ) + + if args.backend in ["sglang", "vllm", "lmdeploy"]: + api_url = ( + f"{args.base_url}/v1/chat/completions" + if args.base_url + else f"http://{args.host}:{args.port}/v1/chat/completions" + ) + base_url = ( + f"http://{args.host}:{args.port}" if args.base_url is None else args.base_url + ) + + # Get model name + if args.model is None: + if args.backend == "truss": + print( + "Please provide a model with `--model` when using truss backend. e.g. --model meta-llama/Llama-3.1-8B-Instruct" + ) + sys.exit(1) + try: + response = requests.get(model_url) + model_list = response.json().get("data", []) + args.model = model_list[0]["id"] if model_list else None + except Exception as e: + print(f"Failed to fetch model from {model_url}. Error: {e}") + print( + "Please specify the correct host and port using `--host` and `--port`." + ) + sys.exit(1) + + if args.model is None: + print("No model specified or found. Please provide a model using `--model`.") + sys.exit(1) + + # Dataset compatibility check + if args.enable_multiturn: + # TODO: Support multiturn for random + if args.dataset_name not in ["sharegpt", "ultrachat", "loogle", "nextqa"]: + print( + "Multiturn conversation is only supported for sharegpt, ultrachat, loogle, and nextqa datasets." + ) + sys.exit(1) + + if args.enable_shared_prefix: + if args.dataset_name not in ["loogle", "nextqa"]: + print("Shared prefix is only supported for loogle and nextqa datasets.") + sys.exit(1) + + print(f"{args}\n") + + # Read dataset + backend = args.backend + model_id = args.model + tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model + + tokenizer = get_tokenizer(tokenizer_id) + + input_requests = get_dataset(args, tokenizer) + + return asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + base_url=base_url, + model_id=model_id, + tokenizer=tokenizer, + input_requests=input_requests, + request_rate=args.request_rate, + max_concurrency=args.max_concurrency, + disable_tqdm=args.disable_tqdm, + lora_name=args.lora_name, + extra_request_body=extra_request_body, + profile=args.profile, + enable_shared_prefix=args.enable_shared_prefix, + ) + ) + + +if __name__ == "__main__": + parser = ArgumentParser(description="Benchmark the online serving throughput.") + parser.add_argument( + "--backend", + type=str, + choices=list(ASYNC_REQUEST_FUNCS.keys()), + default="sglang", + help="Must specify a backend, depending on the LLM Inference Engine.", + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + parser.add_argument( + "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." + ) + parser.add_argument( + "--port", + type=int, + help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=[ + "sharegpt", + "random", + "generated-shared-prefix", + "ultrachat", + "loogle", + "nextqa", + ], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument( + "--dataset-path", type=str, default="", help="Path to the dataset." + ) + parser.add_argument( + "--model", + type=str, + help="Name or path of the model. If not set, the default model will request /v1/models for conf.", + ) + parser.add_argument( + "--tokenizer", + type=str, + help="Name or path of the tokenizer. If not set, using the model conf.", + ) + parser.add_argument( + "--chat-template", + type=str, + help="The buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1000, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--fixed-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output length from the dataset.", + ) + parser.add_argument( + "--sharegpt-context-len", + type=int, + default=None, + help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=1024, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + default=1024, + type=int, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=0.0, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " + "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", + ) + parser.add_argument( + "--max-concurrency", + type=int, + default=None, + help="Maximum number of concurrent requests. This can be used " + "to help simulate an environment where a higher level component " + "is enforcing a maximum number of concurrent requests. While the " + "--request-rate argument controls the rate at which requests are " + "initiated, this argument will control how many are actually allowed " + "to execute at a time. This means that when used in combination, the " + "actual request rate may be lower than specified with --request-rate, " + "if the server is not processing requests fast enough to keep up.", + ) + parser.add_argument( + "--multi", + action="store_true", + help="Use request rate range rather than single value.", + ) + parser.add_argument( + "--request-rate-range", + type=str, + default="2,34,2", + help="Range of request rates in the format start,stop,step. Default is 2,34,2. It also supports a list of request rates, requiring the parameters to not equal three.", + ) + parser.add_argument("--output-file", type=str, help="Output JSONL file name.") + parser.add_argument( + "--enable-multiturn", + action="store_true", + help="Enable multiturn chat for online serving benchmarking. " + "This option is effective on the following datasets: " + "sharegpt, ultrachat, loogle, nextqa", + ) + parser.add_argument( + "--enable-shared-prefix", + action="store_true", + help="Enable shared prefix for online serving benchmarking. " + "This option is effective on the following datasets: " + "loogle, nextqa", + ) + + parser.add_argument( + "--disable-shuffle", + action="store_true", + help="Disable shuffling datasets. This is useful to generate stable output " + "in benchmarking", + ) + parser.add_argument( + "--disable-tqdm", + action="store_true", + help="Specify to disable tqdm progress bar.", + ) + parser.add_argument( + "--disable-stream", + action="store_true", + help="Disable streaming mode.", + ) + parser.add_argument( + "--return-logprob", + action="store_true", + help="Return logprob.", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--disable-ignore-eos", + action="store_true", + help="Disable ignoring EOS.", + ) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + parser.add_argument( + "--apply-chat-template", + action="store_true", + help="Apply chat template", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "SGLANG_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--lora-name", + type=str, + default=None, + help="The name of LoRA adapter", + ) + + group = parser.add_argument_group("generated-shared-prefix dataset arguments") + group.add_argument( + "--gsp-num-groups", + type=int, + default=64, + help="Number of system prompt groups for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-prompts-per-group", + type=int, + default=16, + help="Number of prompts per system prompt group for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-system-prompt-len", + type=int, + default=2048, + help="Target length in tokens for system prompts in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-question-len", + type=int, + default=128, + help="Target length in tokens for questions in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-output-len", + type=int, + default=256, + help="Target length in tokens for outputs in generated-shared-prefix dataset", + ) + # videos specific + parser.add_argument( + "--max-frames", + type=int, + default=sys.maxsize, + help="The maximum number of frames to extract from each video. " + "This option is specific to the nextqa dataset (video benchmark). ", + ) + args = parser.parse_args() + + if args.enable_multiturn and args.enable_shared_prefix: + parser.error( + "--enable-multiturn and --enable-shared-prefix cannot be set at the same time." + ) + + run_benchmark(args) diff --git a/benchmark/hicache/data_processing.py b/benchmark/hicache/data_processing.py new file mode 100644 index 000000000..8f72a0d95 --- /dev/null +++ b/benchmark/hicache/data_processing.py @@ -0,0 +1,590 @@ +import json +import os +import pickle +import random +from pathlib import Path +from typing import List, Optional, Tuple, Union + +import numpy as np +from nextqa import NExTQALoader + +# from nextqa.video import , VideoPrompt +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" + +from sglang.bench_serving import ( + download_and_cache_file, + gen_prompt, + get_gen_prefix_cache_path, +) +from sglang.lang.chat_template import get_chat_template, get_chat_template_by_model_path +from sglang.srt.entrypoints.openai.protocol import ChatCompletionMessageContentPart +from sglang.utils import encode_video_base64 + +# type of content fields, can be only prompts or with images/videos +MsgContent = Union[str, List[ChatCompletionMessageContentPart]] + +# A list of all the conversations. Each conversation is a list of +# tuples. If multiturn is not enabled, the length of list is 1, +# containing only the first Q&A pair. +# For the shared prefix workload (synthetic, loogle, nextqa), it +# is a list of conversations sharing the same prefix (synthetic, +# doc, video) +SampleOutput = List[List[Tuple[MsgContent, int, int]]] + + +def common_filter_chat( + num_requests: int, + new_dataset: List, + tokenizer: PreTrainedTokenizerBase, + min_prompt_len: Optional[int], + min_output_len: Optional[int], + max_prompt_len: Optional[int], + max_output_len: Optional[int], + fixed_output_len: Optional[int], +) -> SampleOutput: + # Filter out sequences that are too long or too short + filtered_dataset: SampleOutput = [] + l = 0 + input_tokens = 0 + output_tokens = 0 + while l < num_requests: + for i in range(len(new_dataset)): + if l == num_requests: + break + processed = [] + for j in new_dataset[i]: + # Tokenize the prompts and completions. + prompt = j[0] + prompt_token_ids = tokenizer.encode(prompt) + prompt_len = len(prompt_token_ids) + + completion = j[1] + completion_token_ids = tokenizer.encode(completion) + output_len = ( + len(completion_token_ids) + if fixed_output_len is None + else fixed_output_len + ) + if ( + min_prompt_len is not None + and prompt_len < min_prompt_len + or min_output_len is not None + and output_len < min_output_len + or max_prompt_len is not None + and prompt_len > max_prompt_len + or max_output_len is not None + and output_len > max_output_len + ): + # Prune too short sequences. + continue + input_tokens += prompt_len + output_tokens += output_len + processed.append((prompt, prompt_len, output_len)) + if len(processed) != 0: + filtered_dataset.append(processed) + l += 1 + + print(f"#Input tokens: {input_tokens}") + print(f"#Output tokens: {output_tokens}") + return filtered_dataset + + +def sample_sharegpt_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + disable_shuffle: bool = False, + enable_multiturn: bool = True, + fixed_output_len: Optional[int] = None, +) -> SampleOutput: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Download sharegpt if necessary + if not os.path.isfile(dataset_path): + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + dataset = [data for data in dataset if len(data["conversations"]) >= 2] + + # Keep one conversation in one list + new_dataset = [] + for data in dataset: + if len(data["conversations"]) % 2 != 0: + continue + if data["conversations"][0]["from"] != "human": + continue + chat = [] + total_len = 2 + if enable_multiturn: + total_len = len(data["conversations"]) + for i in range(0, total_len, 2): + # One user One Assistant + chat.append( + ( + data["conversations"][i]["value"], + data["conversations"][i + 1]["value"], + ) + ) + new_dataset.append(chat) + + if not disable_shuffle: + # Shuffle the dataset. + random.shuffle(new_dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: SampleOutput = common_filter_chat( + num_requests, new_dataset, tokenizer, 4, 4, None, None, fixed_output_len + ) + return filtered_dataset + + +def sample_ultrachat_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + disable_shuffle: bool = False, + enable_multiturn: bool = True, + fixed_output_len: Optional[int] = None, +) -> SampleOutput: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Load the dataset + dataset = [] + with open(dataset_path) as f: + while True: + line = f.readline() + if not line: + break + dataset.append(json.loads(line)) + + # Filter out the conversations with less than 2 turns. + dataset = [data for data in dataset if len(data["data"]) >= 2] + + # Keep one conversation in one list + new_dataset = [] + for data in dataset: + if len(data["data"]) % 2 != 0: + continue + chat = [] + total_len = 2 + if enable_multiturn: + total_len = len(data["data"]) + for i in range(0, total_len, 2): + # One user One Assistant + chat.append((data["data"][i], data["data"][i + 1])) + new_dataset.append(chat) + + # Shuffle the dataset. + if not disable_shuffle: + random.shuffle(new_dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: SampleOutput = common_filter_chat( + num_requests, new_dataset, tokenizer, 4, 4, None, None, fixed_output_len + ) + return filtered_dataset + + +def sample_loogle_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + disable_shuffle: bool = False, + enable_multiturn: bool = True, + enable_shared_prefix: bool = False, + fixed_output_len: Optional[int] = None, +) -> SampleOutput: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Load the dataset + dataset = [] + with open(dataset_path) as f: + while True: + line = f.readline() + if not line: + break + dataset.append(json.loads(line)) + + # Keep one conversation in one list + new_dataset = [] + # TODO: Add shared prefix support for loogle + # NOTE: Now we preprocess it only for chat + for data in dataset: + chat = [] + if ( + "qa_pairs" not in data + or data["qa_pairs"] == "none" + or len(data["qa_pairs"]) == 0 + ): + # If Q is none (for summarization), + # We add a question for summarization + # And keep the summary up to 1024 words + chat.append( + ( + "Input: " + + data["input"] + + " Question: " + + "Please summarize the input", + data["input"][:1024], + ) + ) + new_dataset.append(chat) + else: + qa_pairs = eval(data["qa_pairs"]) + for i, qa in enumerate(qa_pairs): + if i == 0 or enable_shared_prefix: + # Combine input with the first Q + chat.append( + ("Input: " + data["input"] + " Question: " + qa["Q"], qa["A"]) + ) + elif enable_multiturn: + chat.append((qa["Q"], qa["A"])) + + new_dataset.append(chat) + + # Shuffle the dataset. + if not disable_shuffle: + random.shuffle(new_dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: SampleOutput = common_filter_chat( + num_requests, new_dataset, tokenizer, 4, None, None, None, fixed_output_len + ) + return filtered_dataset + + +def sample_nextqa_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + max_frames: int, # Specific for video + model_path: str, + disable_shuffle: bool = False, + enable_multiturn: bool = True, # No multiturn support for now + backend: str = "sglang-oai", + chat_template_name: Optional[str] = None, + fixed_output_len: Optional[int] = None, +) -> SampleOutput: + """ + Example of messages: + message = { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": base64_data}}, + {"type": "text", "text": video.prompt}, + ], + } + """ + + if fixed_output_len is None: + fixed_output_len = 4096 + + # TODO: Check for multiturn + dataset = NExTQALoader(video_dir=dataset_path, max_frames=max_frames) + new_dataset = [] + for v in dataset: + new_dataset.append(v) + + if not disable_shuffle: + random.shuffle(new_dataset) + + # TODO: prompt len can get from server side + filtered_dataset = [] + l = 0 + while l < num_requests: + for i in range(len(new_dataset)): + if l == num_requests: + break + + video = new_dataset[i] + + # text prompt + prompt = video.prompt + + # NOTE: Chat Template is a must for video benchmark because we have to + # add special image token for later expansion + if backend == "sglang" or backend == "sglang-native": + if "chat_template" in tokenizer.init_kwargs: + chat_template = get_chat_template(tokenizer.get_chat_template()) + elif chat_template_name is not None: + chat_template = get_chat_template(chat_template_name) + else: + chat_template = get_chat_template_by_model_path(model_path) + prompt = chat_template.image_token + prompt + + prompt_token_ids = tokenizer(prompt).input_ids + prompt_len = len(prompt_token_ids) + output_len = fixed_output_len # max output len, not real output len + + # video input + base64_data = encode_video_base64(video.path, video.num_frames) + + # NOTE: This will be replaced by the expanded length from the server + prompt_len += video.num_frames + + # add to content + content = [ + {"type": "image_url", "image_url": {"url": base64_data}}, + {"type": "text", "text": prompt}, + ] + + filtered_dataset.append([(content, prompt_len, output_len)]) + l += 1 + return filtered_dataset + + +def sample_random_requests( + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + dataset_path: str, + disable_shuffle: bool = False, +) -> SampleOutput: + + input_lens = np.random.randint( + max(int(input_len * range_ratio), 1), + input_len + 1, + size=num_prompts, + ) + output_lens = np.random.randint( + int(output_len * range_ratio), + output_len + 1, + size=num_prompts, + ) + + if True: + # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens + + # Download sharegpt if necessary + if not os.path.isfile(dataset_path): + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + dataset = [data for data in dataset if len(data["conversations"]) >= 2] + # Only keep the first two turns of each conversation. + dataset = [ + (data["conversations"][0]["value"], data["conversations"][1]["value"]) + for data in dataset + ] + + if not disable_shuffle: + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + input_requests: SampleOutput = [] + for data in dataset: + i = len(input_requests) + if i == num_prompts: + break + + # Tokenize the prompts and completions. + prompt = data[0] + prompt_token_ids = tokenizer.encode(prompt) + prompt_len = len(prompt_token_ids) + + # Skip empty prompt + if prompt_len == 0: + continue + + if prompt_len > input_lens[i]: + input_ids = prompt_token_ids[: input_lens[i]] + else: + ratio = (input_lens[i] + prompt_len - 1) // prompt_len + input_ids = (prompt_token_ids * ratio)[: input_lens[i]] + prompt = tokenizer.decode(input_ids) + input_requests.append([(prompt, int(input_lens[i]), int(output_lens[i]))]) + else: + # Sample token ids from random integers. This can cause some NaN issues. + offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts) + input_requests = [] + for i in range(num_prompts): + prompt = tokenizer.decode( + [ + (offsets[i] + i + j) % tokenizer.vocab_size + for j in range(input_lens[i]) + ] + ) + input_requests.append([(prompt, int(input_lens[i]), int(output_lens[i]))]) + + print(f"#Input tokens: {np.sum(input_lens)}") + print(f"#Output tokens: {np.sum(output_lens)}") + return input_requests + + +def gen_prompt(tokenizer, token_num): + """Generate a random prompt of specified token length using tokenizer vocabulary.""" + all_available_tokens = list(tokenizer.get_vocab().values()) + selected_tokens = random.choices(all_available_tokens, k=token_num) + return tokenizer.decode(selected_tokens) + + +def get_gen_prefix_cache_path(args, tokenizer): + """Create cache directory under ~/.cache/sglang/benchmark""" + cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" + + # Create a unique cache filename based on the generation parameters + cache_key = ( + f"gsp_prefix_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_" + f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_output_len}_" + f"{tokenizer.__class__.__name__}.pkl" + ) + return cache_dir / cache_key + + +def sample_generated_shared_prefix_requests( + num_groups: int, + prompts_per_group: int, + system_prompt_len: int, + question_len: int, + output_len: int, + tokenizer: PreTrainedTokenizerBase, + args, + disable_shuffle: bool = False, +) -> SampleOutput: + """Generate benchmark requests with shared system prompts using random tokens and caching.""" + cache_path = get_gen_prefix_cache_path(args, tokenizer) + + # Try to load from cache first + if cache_path.exists(): + print(f"\nLoading cached generated input data from {cache_path}") + with open(cache_path, "rb") as f: + return pickle.load(f) + + print("\nGenerating new input data...") + + # Generate system prompts for each group + system_prompts = [] + for _ in range(num_groups): + system_prompt = gen_prompt(tokenizer, system_prompt_len) + system_prompts.append(system_prompt) + + # Generate questions + questions = [] + for _ in range(num_groups * prompts_per_group): + question = gen_prompt(tokenizer, question_len) + questions.append(question) + + # Combine system prompts with questions + input_requests = [] + total_input_tokens = 0 + total_output_tokens = 0 + + for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): + system_prompt = system_prompts[group_idx] + input_requests.append([]) + for prompt_idx in tqdm( + range(prompts_per_group), desc="Generating questions", leave=False + ): + question = questions[group_idx * prompts_per_group + prompt_idx] + full_prompt = f"{system_prompt}\n\n{question}" + prompt_len = len(tokenizer.encode(full_prompt)) + input_requests[-1].append((full_prompt, prompt_len, output_len)) + total_input_tokens += prompt_len + total_output_tokens += output_len + + if not disable_shuffle: + # Shuffle questions + random.shuffle(input_requests) + + # Print statistics + print(f"\nGenerated shared prefix dataset statistics:") + print(f"Number of groups: {num_groups}") + print(f"Prompts per group: {prompts_per_group}") + print(f"Total prompts: {len(input_requests) * prompts_per_group}") + print(f"Total input tokens: {total_input_tokens}") + print(f"Total output tokens: {total_output_tokens}") + print( + f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens" + ) + print( + f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n" + ) + + # Save to cache + cache_path.parent.mkdir(parents=True, exist_ok=True) + print(f"Caching generated input data to {cache_path}") + with open(cache_path, "wb") as f: + pickle.dump(input_requests, f) + + return input_requests + + +def get_dataset(args, tokenizer): + if args.dataset_name == "sharegpt": + input_requests = sample_sharegpt_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + disable_shuffle=args.disable_shuffle, + enable_multiturn=args.enable_multiturn, + fixed_output_len=args.fixed_output_len, + ) + elif args.dataset_name == "ultrachat": + input_requests = sample_ultrachat_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + disable_shuffle=args.disable_shuffle, + enable_multiturn=args.enable_multiturn, + fixed_output_len=args.fixed_output_len, + ) + elif args.dataset_name == "loogle": + input_requests = sample_loogle_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + disable_shuffle=args.disable_shuffle, + enable_multiturn=args.enable_multiturn, + enable_shared_prefix=args.enable_shared_prefix, + fixed_output_len=args.fixed_output_len, + ) + elif args.dataset_name == "nextqa": + input_requests = sample_nextqa_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + max_frames=args.max_frames, + model_path=args.model, + disable_shuffle=args.disable_shuffle, + enable_multiturn=args.enable_multiturn, + backend=args.backend, + chat_template_name=args.chat_template, + fixed_output_len=args.fixed_output_len, + ) + elif args.dataset_name == "random": + input_requests = sample_random_requests( + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + dataset_path=args.dataset_path, + ) + elif args.dataset_name == "generated-shared-prefix": + input_requests = sample_generated_shared_prefix_requests( + num_groups=args.gsp_num_groups, + prompts_per_group=args.gsp_prompts_per_group, + system_prompt_len=args.gsp_system_prompt_len, + question_len=args.gsp_question_len, + output_len=args.gsp_output_len, + args=args, + tokenizer=tokenizer, + ) + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + return input_requests diff --git a/benchmark/hicache/download.sh b/benchmark/hicache/download.sh new file mode 100755 index 000000000..340bcc3bc --- /dev/null +++ b/benchmark/hicache/download.sh @@ -0,0 +1,66 @@ +#!/usr/bin/bash + +# The usage function +usage() { + echo "Usage: $0 {sharegpt|ultragpt|loogle|nextqa|all}" + exit 1 +} + +# The download function +download() { + case "$1" in + sharegpt) + echo $1 + wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json + ;; + ultragpt) + echo $1 + # Questions about the world + wget https://cloud.tsinghua.edu.cn/seafhttp/files/be1d7b87-22ca-449e-a6a7-c61d1ea7e010/ultrachat_release_230407.json + # Writing and Creation + wget https://cloud.tsinghua.edu.cn/seafhttp/files/61742d2a-25e2-4d08-b2b9-15f47ae50ace/ultrachat_material_release_230417.json + wget https://cloud.tsinghua.edu.cn/seafhttp/files/f71f6aa6-d346-4b16-85b7-8502efa3d608/ultrachat_material_release_230412.json + # External materials + wget https://cloud.tsinghua.edu.cn/seafhttp/files/42d22e28-e899-4975-a70f-5eda163e265d/ultrachat_existent_material_release_230420.json.gz + gunzip ultrachat_existent_material_release_230420.json.gz + ;; + loogle) + echo $1 + git lfs install + git clone git@hf.co:datasets/bigainlco/LooGLE + unzip LooGLE/data.zip + ;; + nextqa) + echo $1 + git lfs install + git clone https://huggingface.co/datasets/lmms-lab/NExTQA + unzip NExTQA/videos.zip + ;; + *) + usage + exit 1 + ;; + esac +} + +# Arg check +if [ "$#" -ne 1 ]; then + usage +fi + +# Invoke + +case "$1" in + sharegpt|ultragpt|loogle|nextqa) + download "$1" + ;; + all) + download sharegpt + download ultragpt + download loogle + download nextqa + ;; + *) + usage + ;; +esac diff --git a/benchmark/hicache/nextqa.py b/benchmark/hicache/nextqa.py new file mode 100644 index 000000000..4db6caa1c --- /dev/null +++ b/benchmark/hicache/nextqa.py @@ -0,0 +1,159 @@ +import os +import sys +from typing import List + +import av +from datasets import load_dataset + + +def find_video_files(video_dir) -> List[str]: + if os.path.isfile(video_dir): + return [video_dir] + + video_files = [] + for root, dirs, files in os.walk(video_dir): + for file in files: + if file.endswith((".mp4", ".avi", ".mov")): + video_files.append(os.path.join(root, file)) + # if file is dir + elif os.path.isdir(file): + video_files.extend(find_video_files(file)) + return video_files + + +def video_frames(video_path, max_frames) -> int: + container = av.open(video_path) + total_frames = container.streams.video[0].frames + return min(total_frames, max_frames) + + +class Video: + def __init__(self, video_path, num_frames): + self.path = video_path + self.num_frames = num_frames + + def __str__(self): + return f"Video({self.path}, {self.num_frames})" + + def __iter__(self): + return iter((self.path, self.num_frames)) + + +class VideoPrompt(Video): + def __init__(self, video_path, num_frames, prompt): + super().__init__(video_path, num_frames) + self.prompt = prompt + + def __str__(self): + return f"VideoPrompt({self.path}, {self.num_frames}, {self.prompt})" + + def __iter__(self): + return iter((self.path, self.num_frames, self.prompt)) + + +class VideoLoader: + pass + + +class VideoFileLoader(VideoLoader): + """ + Load all the videos in a directory + """ + + def __init__(self, video_dir, batch_size=1, max_frames=sys.maxsize): + super().__init__() + self.video_dir = video_dir + self.video_files = find_video_files(video_dir) + self.batch_size = batch_size + self.max_frames = max_frames + print(f"batch_size: {batch_size}, max_frames: {max_frames}") + + def __iter__(self): # (file, number of frames) + if self.batch_size == 1: + for video_file in self.video_files: + yield Video(video_file, video_frames(video_file, self.max_frames)) + else: + batch = [] + for video_file in self.video_files: + video = Video(video_file, video_frames(video_file, self.max_frames)) + batch.append(video) + if len(batch) == self.batch_size: + yield batch + batch = [] + + +class NExTQALoader(VideoLoader): + """ + Load vdideos and prompts from NExT dataset + set: train, test or validation + """ + + def __init__( + self, video_dir, batch_size=1, max_frames=sys.maxsize, dset="test", task="OE" + ): + """ + task: 'MV' or 'OE' + """ + super().__init__() + self.task = task + print(f"Loading the {dset} data of {task} from lmms-lab/NExTQA") + self.ds = load_dataset("lmms-lab/NExTQA", task) + self.ds = self.ds[dset] + + # self.n = ds.num_rows + self.video_dir = video_dir + self.video_files = find_video_files(video_dir) + self.video_to_path = dict() + for video_file in self.video_files: + video_id = video_file.split("/")[-1].split(".")[0] + self.video_to_path[video_id] = video_file + + self.batch_size = batch_size + self.max_frames = max_frames + + def get_video_prompt(self, entry, max_frames) -> VideoPrompt: + # Get video + video_id = entry["video"] + video_path = self.video_to_path[video_id] + assert os.path.exists(video_path), f"Video not found: {video_path}" + num_frames = min(entry["frame_count"], max_frames) + video = Video(video_path, num_frames) + prompt = entry["question"] + "?" + if self.task == "MC": # add choices + prompt += f' a0: {entry["a0"]}, a1: {entry["a1"]}, a2: {entry["a2"]}, a3: {entry["a3"]}' + return VideoPrompt(video_path, num_frames, prompt) + + def __iter__(self): + if self.batch_size == 1: + for entry in self.ds: + yield self.get_video_prompt(entry, self.max_frames) + else: + batch = [] + for entry in self.ds: + video = self.get_video_prompt(entry, self.max_frames) + batch.append(video) + if len(batch) == self.batch_size: + yield batch + batch = [] + + +# main +if __name__ == "__main__": + video_dir = "./videos" + # video_loader = VideoFileLoader(video_dir, batch_size=16) + # for batch in video_loader: + # print(f"Number of videos in batch: {len(batch)}") + # for video_file, num_frames in batch: + # print(f"Video: {video_file} number of frames: {num_frames}") + + video_loader = NExTQALoader(video_dir, batch_size=16, dset="test", task="OE") + for batch in video_loader: + print(f"Number of videos in batch: {len(batch)}") + for video_file, num_frames, prompt in batch: + print( + f"Video: {video_file} number of frames: {num_frames}, prompt: {prompt}" + ) + # break + # for video_file, prompt in batch: + # print(f"Video: {video_file} prompt: {prompt}") + # break diff --git a/benchmark/json_decode_regex/README.md b/benchmark/json_decode_regex/README.md new file mode 100644 index 000000000..26acabae9 --- /dev/null +++ b/benchmark/json_decode_regex/README.md @@ -0,0 +1,60 @@ +## Run benchmark + +### Build dataset +``` +pip install wikipedia +python3 build_dataset.py +``` + +### Dependencies + +``` +llama_cpp_python 0.2.19 +guidance 0.1.10 +vllm 0.2.5 +outlines 0.0.22 +``` + +### Benchmark sglang + +Run Llama-7B + +``` +python3 -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +Run Mixtral-8x7B + +``` +python3 -m sglang.launch_server --model-path mistralai/Mixtral-8x7B-Instruct-v0.1 --port 30000 --tp-size 8 +``` + +Benchmark + +``` +python3 bench_sglang.py --num-questions 10 +``` + + +### Benchmark Outlines + vLLM + +Run Llama-7B + +``` +python3 -m outlines.serve.serve --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +Benchmark + +``` +python3 bench_other.py --backend outlines --num-questions 10 +``` + + +### Benchmark guidance + +Run Llama-7B and benchmark + +``` +python3 bench_other.py --backend guidance --num-questions 10 --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` diff --git a/benchmark/json_decode_regex/bench_other.py b/benchmark/json_decode_regex/bench_other.py new file mode 100644 index 000000000..87051ea82 --- /dev/null +++ b/benchmark/json_decode_regex/bench_other.py @@ -0,0 +1,98 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.lang.ir import REGEX_FLOAT, REGEX_INT, REGEX_STR +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +REGEX_LIST = r"\[(" + REGEX_STR + ", )*" + REGEX_STR + r"\]" + + +# fmt: off +def json_decode(document, generate): + s = "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += "{\n" + s += ' "name": ' + s += generate(s, max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "country": ' + s += generate(s, max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "latitude": ' + s += generate(s, max_tokens=8, regex=REGEX_FLOAT + ",") + "\n" + s += ' "population": ' + s += generate(s, max_tokens=8, regex=REGEX_INT + ",") + "\n" + s += ' "top 3 landmarks": ' + s += generate(s, max_tokens=24, regex=REGEX_LIST) + "\n" + s += "}\n" + + return s +# fmt: on + + +def main(args): + lines = read_jsonl(args.data_path) + arguments = [] + for i in range(len(lines[: args.num_questions])): + arguments.append( + { + "document": lines[i]["document"], + } + ) + states = [None] * len(arguments) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + def get_one_answer(i): + states[i] = json_decode(generate=call_generate, **arguments[i]) + + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(arguments))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + rets = list( + tqdm( + executor.map(get_one_answer, list(range(len(arguments)))), + total=len(arguments), + ) + ) + for _ in rets: + pass + + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "json_decode_regex", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=20) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/json_decode_regex/bench_sglang.py b/benchmark/json_decode_regex/bench_sglang.py new file mode 100644 index 000000000..9aab11e43 --- /dev/null +++ b/benchmark/json_decode_regex/bench_sglang.py @@ -0,0 +1,101 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.lang.ir import REGEX_FLOAT, REGEX_INT, REGEX_STR +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +REGEX_LIST = r"\[(" + REGEX_STR + ", )*" + REGEX_STR + r"\]" + +# fmt: off +@sgl.function +def json_warm_up(s): + s += "The information about Hogwarts is in the following JSON format.\n" + with s.var_scope("json_output"): + s += "{\n" + s += ' "name": ' + sgl.gen("name", max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "country": ' + sgl.gen("country", max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "latitude": ' + sgl.gen("latitude", max_tokens=8, regex=REGEX_FLOAT + ",") + "\n" + s += ' "population": ' + sgl.gen("population", max_tokens=8, regex=REGEX_INT + ",") + "\n" + s += ' "top 3 landmarks": ' + sgl.gen( "landmarks", max_tokens=24, regex=REGEX_LIST) + "\n" + s += "}\n" + print(f'The warmp up json result is:\n{s["json_output"]}') +# fmt: on + +# fmt: off +@sgl.function +def json_decode(s, document): + s += "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + with s.var_scope("json_output"): + s += "{\n" + s += ' "name": ' + sgl.gen("name", max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "country": ' + sgl.gen("country", max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "latitude": ' + sgl.gen("latitude", max_tokens=8, regex=REGEX_FLOAT + ",") + "\n" + s += ' "population": ' + sgl.gen("population", max_tokens=8, regex=REGEX_INT + ",") + "\n" + s += ' "top 3 landmarks": ' + sgl.gen( "landmarks", max_tokens=24, regex=REGEX_LIST) + "\n" + s += "}\n" +# fmt: on + + +def main(args): + lines = read_jsonl(args.data_path) + lines = list(lines) + arguments = [] + for i in range(len(lines[: args.num_questions])): + arguments.append( + { + "document": lines[i]["document"], + } + ) + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Warm up + json_warm_up.run().sync() + + # Run requests + tic = time.perf_counter() + states = json_decode.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(f"tmp_{args.backend}_json_results.txt", "w") as fout: + for state in states: + fout.write(state["json_output"] + "\n") + + with open(args.result_file, "a") as fout: + value = { + "task": "json_decode_regex", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=20) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/json_decode_regex/build_dataset.py b/benchmark/json_decode_regex/build_dataset.py new file mode 100644 index 000000000..1396e5ede --- /dev/null +++ b/benchmark/json_decode_regex/build_dataset.py @@ -0,0 +1,58 @@ +import json + +import transformers +import wikipedia + +model_path = "meta-llama/Llama-2-7b-chat-hf" +t = transformers.AutoTokenizer.from_pretrained(model_path) +city_names = [ + "los angles", + "london", + "tokyo", + "beijing", + "singapore", + "paris", + "dubai", + "sydney", + "moscow", + "rome", + "toronto", + "rio de janeiro", + "istanbul", + "berlin", + "auckland", + "buenos aires", + "mexico city", + "mumbai", + "seoul", + "bangkok", + "cairo", + "athens", + "jerusalem", +] + + +def get_content(city_name): + content = str(wikipedia.page(city_name).content) + content = content.replace("\n\n", "\n") + + tokens = t.encode(content) + + expected_tokens = 3000 + truncate_len = int((expected_tokens / len(tokens)) * len(content)) + truncate_content = content[:truncate_len] + truncate_tokens = t.encode(truncate_content) + + # Count token + print( + f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}" + ) + + return truncate_content + + +if __name__ == "__main__": + with open("questions.jsonl", "w") as fout: + for city_name in city_names: + truncate_content = get_content(city_name) + fout.write(json.dumps({"document": truncate_content}) + "\n") diff --git a/benchmark/json_jump_forward/README.md b/benchmark/json_jump_forward/README.md new file mode 100644 index 000000000..38fb67e89 --- /dev/null +++ b/benchmark/json_jump_forward/README.md @@ -0,0 +1,88 @@ +## Run benchmark + +### Dependencies + +``` +llama_cpp_python 0.2.38 +guidance 0.1.10 +vllm 0.2.7 +outlines 0.0.25 +``` + +### Build dataset + +When benchmarking long document information retrieval, run the following command to build the dataset: + +```bash +pip install wikipedia +python3 build_dataset.py +``` + +### Benchmark sglang + +Run Llama-7B + +```bash +python3 -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +Benchmark Character Generation + +```bash +python3 bench_sglang.py --mode character +``` + +Benchmark City Information Retrieval + +```bash +python3 bench_sglang.py --mode city +``` + + +### Benchmark Outlines + vLLM + +Run Llama-7B + +```bash +python3 -m outlines.serve.serve --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +Benchmark Character Generation + +```bash +python3 bench_other.py --mode character --backend outlines +``` + +Benchmark City Information Retrieval + +```bash +python3 bench_other.py --mode city --backend outlines +``` + +### Benchmark guidance + +Run Llama-7B and benchmark character generation + +```bash +python3 bench_other.py --mode character --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +Run Llama-7B and benchmark city information retrieval + +```bash +python3 bench_other.py --mode city --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +Run Llama-7B and benchmark character generation + +``` +python3 bench_other.py --mode character --backend lmql --parallel 1 +``` + +Run Llama-7B and benchmark city information retrieval + +``` +python3 bench_other.py --mode city --backend lmql --parallel 1 +``` diff --git a/benchmark/json_jump_forward/bench_other.py b/benchmark/json_jump_forward/bench_other.py new file mode 100644 index 000000000..a64e950d7 --- /dev/null +++ b/benchmark/json_jump_forward/bench_other.py @@ -0,0 +1,288 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +import guidance +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +# there are some FSM bugs with json regex converted from pydantic model +# here use a string regex instead +# regex_string = build_regex_from_object(HarryPoterRole) +character_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "house": "(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)",\n""" + + r""" "blood status": "(Pure-blood|Half-blood|Muggle-born)",\n""" + + r""" "occupation": "(student|teacher|auror|ministry of magic|death eater|order of the phoenix)",\n""" + + r""" "wand": \{\n""" + + r""" "wood": "[\w\d\s]{1,16}",\n""" + + r""" "core": "[\w\d\s]{1,16}",\n""" + + r""" "length": [0-9]{1,2}\.[0-9]{0,2}\n""" + + r""" \},\n""" + + r""" "alive": "(Alive|Deceased)",\n""" + + r""" "patronus": "[\w\d\s]{1,16}",\n""" + + r""" "bogart": "[\w\d\s]{1,16}"\n""" + + r"""\}""" +) + +city_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "country": "[\w\d\s]{1,16}",\n""" + + r""" "latitude": [-+]?[0-9]*\.?[0-9]{0,2},\n""" + + r""" "population": [-+]?[0-9]{1,9},\n""" + + r""" "top 3 landmarks": \["[\w\d\s]{1,16}", "[\w\d\s]{1,16}", "[\w\d\s]{1,16}"\]\n""" + + r"""\}""" +) + +# fmt: off +def character_gen(name, generate): + s = name + " is a character in Harry Potter. Please fill in the following information about this character.\n" + s += generate(s, max_tokens=256, regex=character_regex) + return s +# fmt: on + +# fmt: off +def city_gen(document, generate): + s = "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += generate(s, max_tokens=256, regex=city_regex) + return s +# fmt: on + + +@guidance +def character_maker(lm, name): + regex_str_no_quote = r"[\w\d\s]+" + regex_float = r"[0-9]+\.[0-9]+" + lm += f"""\ + {name} is a character in Harry Potter. Please fill in the following information about this character. + {{ + "name": "{guidance.gen("name", max_tokens=16, regex=regex_str_no_quote)}", + "house": "{guidance.select(options=['Gryffindor', 'Slytherin', 'Ravenclaw', 'Hufflepuff'], name='house')}", + "blood status": "{guidance.select(options=['Pure-blood', 'Half-blood', 'Muggle-born'], name='blood status')}", + "occupation": "{guidance.select(options=['student', 'teacher', 'auror', 'ministry of magic', 'death eater', 'order of the phoenix'], name='occupation')}", + "wand": {{ + "wood": "{guidance.gen("wood", max_tokens=16, regex=regex_str_no_quote)}", + "core": "{guidance.gen('core', max_tokens=16, regex=regex_str_no_quote)}", + "length": {guidance.gen('length', max_tokens=10, regex=regex_float)} + }}, + "alive": "{guidance.select(options=['Alive', 'Deceased'], name='alive')}", + "patronus": "{guidance.gen('patronus', max_tokens=16, regex=regex_str_no_quote)}", + "bogart": "{guidance.gen('bogart', max_tokens=16, regex=regex_str_no_quote)}" + }} + """ + + return lm + + +async def call_generate_lmql( + prompt, temperature, max_tokens, regex, max_len=4096, model=None, **kwargs +): + assert model is not None + import lmql + + @lmql.query(model=model) + async def program(question, max_tokens, regex): + '''lmql + """{question}[ANSWER]""" where len(TOKENS(ANSWER)) < max_tokens and REGEX(ANSWER, regex) + return ANSWER + ''' + + return await program( + question=prompt, + temperature=temperature, + max_tokens=max_tokens, + max_len=max_len, + regex=regex, + **kwargs, + ) + + +@guidance +def city_maker(lm, document): + regex_str_no_quote = r"[\w\d\s]+" + regex_float = r"[0-9]+\.[0-9]+" + lm += f"""\ + Please extract the information of a city from the following wikipedia page. + Page begin. + {document} + Page end. + Here is the name, country, and symbol of the city in JSON format. + {{ + "name": "{guidance.gen("name", max_tokens=16, regex=regex_str_no_quote)}", + "country": "{guidance.gen("country", max_tokens=16, regex=regex_str_no_quote)}", + "latitude": {guidance.gen("latitude", max_tokens=10, regex=regex_float)}, + "population": {guidance.gen("population", max_tokens=10, regex=r"[0-9]+")}, + "top 3 landmarks": [ + "{guidance.gen("landmark1", max_tokens=16, regex=regex_str_no_quote)}", "{guidance.gen("landmark2", max_tokens=16, regex=regex_str_no_quote)}", "{guidance.gen("landmark3", max_tokens=16, regex=regex_str_no_quote)}" + ] + }} + """ + + return lm + + +def bench_character(args): + arguments = [] + with open(args.data_path, "r") as f: + for line in f: + arguments.append({"name": line.strip()}) + arguments = arguments[: args.num_jsons] + + states = [None] * len(arguments) + + # Select backend + if args.backend == "outlines": + call_generate = partial(get_call_generate(args), temperature=0) + + def get_one_answer(i): + states[i] = character_gen(**arguments[i], generate=call_generate) + + elif args.backend == "guidance": + model = guidance.models.LlamaCpp( + args.model_path, + n_gpu_layers=-1, + n_ctx=args.n_ctx, + ) + + def get_one_answer(i): + lm = model + character_maker(**arguments[i]) + states[i] = lm + + elif args.backend == "lmql": + import asyncio + + import lmql + + model = lmql.model(args.model_path, endpoint=f"{args.host}:{args.port}") + call_generate = partial( + call_generate_lmql, + model=model, + max_tokens=256, + regex=character_regex, + ) + + async def get_one_answer_async(i): + states[i] = await call_generate(prompt=arguments[i]["name"], temperature=0) + + else: + raise ValueError(f"Invalid backend: {args.backend}") + + tic = time.perf_counter() + + if args.backend != "lmql": + if args.parallel == 1: + for i in tqdm(range(len(arguments))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + rets = list( + tqdm( + executor.map(get_one_answer, list(range(len(arguments)))), + total=len(arguments), + ) + ) + for _ in rets: + pass + else: + batches = [] + for i in range(0, len(arguments), args.parallel): + batches.append(list(range(i, min(i + args.parallel, len(arguments))))) + loop = asyncio.get_event_loop() + + for bt in tqdm(batches): + loop.run_until_complete( + asyncio.gather(*[get_one_answer_async(i) for i in bt]) + ) + + latency = time.perf_counter() - tic + + return states, latency + + +def bench_city_doc(args): + arguments = [] + for line in read_jsonl(args.data_path): + arguments.append({"document": line["document"]}) + arguments = arguments[: args.num_jsons] + + states = [None] * len(arguments) + + # Select backend + if args.backend == "outlines": + call_generate = partial(get_call_generate(args), temperature=0) + + def get_one_answer(i): + states[i] = city_gen(**arguments[i], generate=call_generate) + + elif args.backend == "guidance": + model = guidance.models.LlamaCpp( + args.model_path, + n_gpu_layers=-1, + n_ctx=args.n_ctx, + ) + + def get_one_answer(i): + lm = model + city_maker(**arguments[i]) + states[i] = lm + + else: + raise ValueError(f"Invalid backend: {args.backend}") + + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(arguments))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + rets = executor.map(get_one_answer, list(range(len(arguments)))) + for _ in rets: + pass + + latency = time.perf_counter() - tic + + return states, latency + + +def main(args): + if args.mode == "character": + args.data_path = "dataset.txt" + states, latency = bench_character(args) + elif args.mode == "city": + args.data_path = "questions.jsonl" + states, latency = bench_city_doc(args) + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}_{args.mode}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "json_jump_forward", + "backend": args.backend, + "latency": round(latency, 3), + "num_jsons": args.num_jsons, + "mode": args.mode, + "parallel": args.parallel, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str) + parser.add_argument("--num-jsons", type=int, default=50) + parser.add_argument( + "--mode", type=str, default="character", choices=["character", "city"] + ) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/json_jump_forward/bench_sglang.py b/benchmark/json_jump_forward/bench_sglang.py new file mode 100644 index 000000000..29f635f75 --- /dev/null +++ b/benchmark/json_jump_forward/bench_sglang.py @@ -0,0 +1,143 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +# there are some FSM bugs with json regex converted from pydantic model +# here use a string regex instead +# regex_string = build_regex_from_object(HarryPoterRole) +character_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "house": "(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)",\n""" + + r""" "blood status": "(Pure-blood|Half-blood|Muggle-born)",\n""" + + r""" "occupation": "(student|teacher|auror|ministry of magic|death eater|order of the phoenix)",\n""" + + r""" "wand": \{\n""" + + r""" "wood": "[\w\d\s]{1,16}",\n""" + + r""" "core": "[\w\d\s]{1,16}",\n""" + + r""" "length": [0-9]{1,2}\.[0-9]{0,2}\n""" + + r""" \},\n""" + + r""" "alive": "(Alive|Deceased)",\n""" + + r""" "patronus": "[\w\d\s]{1,16}",\n""" + + r""" "bogart": "[\w\d\s]{1,16}"\n""" + + r"""\}""" +) + +city_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "country": "[\w\d\s]{1,16}",\n""" + + r""" "latitude": [-+]?[0-9]*\.?[0-9]{0,2},\n""" + + r""" "population": [-+]?[0-9]{1,9},\n""" + + r""" "top 3 landmarks": \["[\w\d\s]{1,16}", "[\w\d\s]{1,16}", "[\w\d\s]{1,16}"\]\n""" + + r"""\}""" +) + +# fmt: off +@sgl.function +def character_gen(s, name): + s += name + " is a character in Harry Potter. Please fill in the following information about this character.\n" + s += sgl.gen("json_output", max_tokens=256, regex=character_regex) +# fmt: on + +# fmt: off +@sgl.function +def city_gen(s, document): + s += "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += sgl.gen("json_output",max_tokens=256, regex=city_regex) +# fmt: on + + +def bench_city_doc(args): + arguments = [] + for line in read_jsonl(args.data_path): + arguments.append({"document": line["document"]}) + arguments = arguments[: args.num_jsons] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.perf_counter() + states = city_gen.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + return states, latency + + +def bench_character(args): + arguments = [] + with open(args.data_path, "r") as f: + for line in f: + arguments.append({"name": line.strip()}) + arguments = arguments[: args.num_jsons] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.perf_counter() + states = character_gen.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + return states, latency + + +def main(args): + if args.mode == "character": + args.data_path = "dataset.txt" + states, latency = bench_character(args) + elif args.mode == "city": + args.data_path = "questions.jsonl" + states, latency = bench_city_doc(args) + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}_{args.mode}.txt", states) + with open(f"{args.backend}_{args.mode}.json", "w") as fout: + for state in states: + fout.write(state["json_output"] + "\n") + + with open(args.result_file, "a") as fout: + value = { + "task": "json_jump_forward", + "backend": args.backend, + "latency": round(latency, 3), + "num_jsons": args.num_jsons, + "mode": args.mode, + "parallel": args.parallel, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str) + parser.add_argument("--num-jsons", type=int, default=50) + parser.add_argument( + "--mode", type=str, default="character", choices=["character", "city"] + ) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/json_jump_forward/build_dataset.py b/benchmark/json_jump_forward/build_dataset.py new file mode 100644 index 000000000..1396e5ede --- /dev/null +++ b/benchmark/json_jump_forward/build_dataset.py @@ -0,0 +1,58 @@ +import json + +import transformers +import wikipedia + +model_path = "meta-llama/Llama-2-7b-chat-hf" +t = transformers.AutoTokenizer.from_pretrained(model_path) +city_names = [ + "los angles", + "london", + "tokyo", + "beijing", + "singapore", + "paris", + "dubai", + "sydney", + "moscow", + "rome", + "toronto", + "rio de janeiro", + "istanbul", + "berlin", + "auckland", + "buenos aires", + "mexico city", + "mumbai", + "seoul", + "bangkok", + "cairo", + "athens", + "jerusalem", +] + + +def get_content(city_name): + content = str(wikipedia.page(city_name).content) + content = content.replace("\n\n", "\n") + + tokens = t.encode(content) + + expected_tokens = 3000 + truncate_len = int((expected_tokens / len(tokens)) * len(content)) + truncate_content = content[:truncate_len] + truncate_tokens = t.encode(truncate_content) + + # Count token + print( + f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}" + ) + + return truncate_content + + +if __name__ == "__main__": + with open("questions.jsonl", "w") as fout: + for city_name in city_names: + truncate_content = get_content(city_name) + fout.write(json.dumps({"document": truncate_content}) + "\n") diff --git a/benchmark/json_jump_forward/dataset.txt b/benchmark/json_jump_forward/dataset.txt new file mode 100644 index 000000000..c12421e5d --- /dev/null +++ b/benchmark/json_jump_forward/dataset.txt @@ -0,0 +1,50 @@ +Harry Potter +Hermione Granger +Ron Weasley +Albus Dumbledore +Severus Snape +Rubeus Hagrid +Draco Malfoy +Ginny Weasley +Fred Weasley +George Weasley +Percy Weasley +Sirius Black +Remus Lupin +Neville Longbottom +Luna Lovegood +Cedric Diggory +Cho Chang +Lord Voldemort +Minerva McGonagall +Filius Flitwick +Dolores Umbridge +Bellatrix Lestrange +Lucius Malfoy +Molly Weasley +Arthur Weasley +Nymphadora Tonks +Dobby +Moaning Myrtle +Peter Pettigrew +Alastor 'Mad-Eye' Moody +Horace Slughorn +Vernon Dursley +Petunia Dursley +Dudley Dursley +Argus Filch +Sybill Trelawney +Gilderoy Lockhart +Fleur Delacour +Viktor Krum +Bill Weasley +Oliver Wood +Cornelius Fudge +Barty Crouch Sr. +Barty Crouch Jr. +Kingsley Shacklebolt +Quirinus Quirrell +Nearly Headless Nick +Aunt Marge +Griphook +Ludo Bagman diff --git a/benchmark/json_schema/README.md b/benchmark/json_schema/README.md new file mode 100644 index 000000000..e6790e1fc --- /dev/null +++ b/benchmark/json_schema/README.md @@ -0,0 +1,15 @@ +## Run benchmark + +### Benchmark sglang + +Run Llama-8b + +```bash +python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --port 30000 +``` + +Benchmark + +```bash +python3 bench_sglang.py +``` diff --git a/benchmark/json_schema/bench_sglang.py b/benchmark/json_schema/bench_sglang.py new file mode 100644 index 000000000..55365ff2e --- /dev/null +++ b/benchmark/json_schema/bench_sglang.py @@ -0,0 +1,146 @@ +import argparse +import json +import time +from typing import List, Tuple + +import jsonschema +from datasets import load_dataset + +import sglang as sgl +from sglang.global_config import global_config +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +@sgl.function +def schema_gen(s, message: Tuple[str, str], json_schema: str): + system, user = message + s += sgl.system(system) + s += sgl.user(user) + s += sgl.assistant( + sgl.gen("json_output", temperature=0, max_tokens=256, json_schema=json_schema) + ) + + +def contains_formats(schema, formats: List[str]): + if isinstance(schema, dict): + if schema.get("format", None) in formats: + return True + for value in schema.values(): + if contains_formats(value, formats): + return True + elif isinstance(schema, list): + for item in schema: + if contains_formats(item, formats): + return True + return False + + +def convert_dataset(path: str): + raw_dataset = load_dataset(path) + dataset = [] + for data in raw_dataset["train"]: + messages = data["prompt"] + schema = data["schema"] + obj = json.loads(schema) + + # skip some corrupted examples + if obj.get("type", None) is None: + continue + + # skip schema with format "email" + # which is not supported by outlines for now + if contains_formats(obj, ["email"]): + continue + + system = messages[0] + user = messages[1] + assert system["role"] == "system", "invalid role" + assert user["role"] == "user", "invalid role" + assert len(messages) == 2, "invalid message length" + message = json.dumps(system["content"]), json.dumps(user["content"]) + dataset.append( + { + "message": message, + "json_schema": schema, + } + ) + + return dataset + + +def bench_schema(args): + arguments = convert_dataset(args.data_path) + + if args.num_jsons < 0 or args.num_jsons > len(arguments): + args.num_jsons = len(arguments) + arguments = arguments[: args.num_jsons] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.perf_counter() + states = schema_gen.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + # Check if the outputs are valid + indexes = [] + for i, state in enumerate(states): + try: + schema = json.loads(arguments[i]["json_schema"]) + obj = json.loads(state["json_output"]) + assert jsonschema.validate(obj, schema) is None + except Exception as e: + print(e) + indexes.append(i) + + return states, latency + + +def main(args): + states, latency = bench_schema(args) + + # Compute accuracy + tokenizer = get_tokenizer( + global_config.default_backend.get_server_info()["tokenizer_path"] + ) + output_jsons = [state["json_output"] for state in states] + num_output_tokens = sum(len(tokenizer.encode(x)) for x in output_jsons) + print(f"Latency: {latency:.3f}") + print(f"Output throughput: {num_output_tokens / latency:.3f} token/s") + print(f"#output tokens: {num_output_tokens}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + with open(f"{args.backend}.jsonl", "w") as fout: + for state in states: + fout.write(state["json_output"] + "\n") + + with open(args.result_file, "a") as fout: + value = { + "task": "json_schema", + "backend": args.backend, + "latency": round(latency, 3), + "num_jsons": args.num_jsons, + "parallel": args.parallel, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="NousResearch/json-mode-eval") + parser.add_argument("--num-jsons", type=int, default=-1) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/kernels/all_reduce/benchmark_mscclpp.py b/benchmark/kernels/all_reduce/benchmark_mscclpp.py new file mode 100644 index 000000000..eebbd00ce --- /dev/null +++ b/benchmark/kernels/all_reduce/benchmark_mscclpp.py @@ -0,0 +1,224 @@ +"""For Now, MSCCL is only supported on TP16 and TP8 case + +export WORLD_SIZE=1 +export RANK=0 +export MASTER_ADDR=127.0.0.1 +export MASTER_PORT=12345 + +torchrun --nproc_per_node gpu \ +--nnodes $WORLD_SIZE \ +--node_rank $RANK \ +--master_addr $MASTER_ADDR \ +--master_port $MASTER_PORT benchmark/kernels/all_reduce/benchmark_mscclpp.py +""" + +import os +from contextlib import nullcontext +from typing import List + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from sglang.srt.distributed import init_distributed_environment +from sglang.srt.distributed.device_communicators.pymscclpp import PyMscclppCommunicator +from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator +from sglang.srt.distributed.parallel_state import ( + get_tensor_model_parallel_group, + graph_capture, + initialize_model_parallel, + set_mscclpp_all_reduce, +) + + +def torch_allreduce(torch_input: torch.Tensor, group: ProcessGroup) -> torch.Tensor: + dist.all_reduce(torch_input, group=group) + return torch_input + + +def msccl_allreduce( + msccl_input: torch.Tensor, msccl_comm: PyMscclppCommunicator +) -> torch.Tensor: + return msccl_comm.all_reduce(msccl_input) + + +def pynccl_allreduce( + msccl_input: torch.Tensor, pynccl_comm: PyNcclCommunicator +) -> torch.Tensor: + pynccl_comm.all_reduce(msccl_input) + return msccl_input + + +def _bench_graph_time(func, inp_randn, warmup_loop=2, graph_loop=10, test_loop=10): + graph_input = inp_randn.clone() + with graph_capture() as graph_capture_context: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=graph_capture_context.stream): + for _ in range(graph_loop): + graph_out = func(graph_input) + + graph.replay() + func_output = graph_out.clone() + + for _ in range(warmup_loop): + graph.replay() + torch.cuda.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + latencies: List[float] = [] + for _ in range(test_loop): + torch.cuda.synchronize() + dist.barrier() + start_event.record() + graph.replay() + end_event.record() + end_event.synchronize() + latencies.append(start_event.elapsed_time(end_event)) + func_cost_us = sum(latencies) / len(latencies) / graph_loop * 1000 + graph.reset() + return func_output, func_cost_us + + +def _bench_eager_time(func, inp_randn, warmup_loop=2, test_loop=10): + eager_input = inp_randn.clone() + eager_output = func(eager_input) + func_output = eager_output.clone() + + for _ in range(warmup_loop): + func(eager_input) + torch.cuda.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start_event.record() + for _ in range(test_loop): + func(eager_input) + end_event.record() + torch.cuda.synchronize() + func_cost_us = start_event.elapsed_time(end_event) / test_loop * 1000 + + return func_output, func_cost_us + + +def get_torch_prof_ctx(do_prof: bool): + ctx = ( + torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + with_stack=True, + ) + if do_prof + else nullcontext() + ) + return ctx + + +def human_readable_size(size, decimal_places=1): + for unit in ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]: + if size < 1024.0 or unit == "PiB": + break + size /= 1024.0 + return f"{size:.{decimal_places}f} {unit}" + + +try: + from tabulate import tabulate +except ImportError: + print("tabulate not installed, skipping table printing") + tabulate = None + + +def print_markdown_table(data): + if tabulate is not None: + print(tabulate(data, headers="keys", tablefmt="github")) + return + headers = data[0].keys() + header_row = "| " + " | ".join(headers) + " |" + separator = "| " + " | ".join(["---"] * len(headers)) + " |" + rows = [] + for item in data: + row = "| " + " | ".join(str(item[key]) for key in headers) + " |" + rows.append(row) + markdown_table = "\n".join([header_row, separator] + rows) + print(markdown_table) + + +if __name__ == "__main__": + import logging + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + force=True, + ) + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + world, world_size = dist.group.WORLD, dist.get_world_size() + rank = dist.get_rank() + torch.cuda.set_device(rank % 8) + device = torch.cuda.current_device() + set_mscclpp_all_reduce(True) + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=rank % 8, + ) + initialize_model_parallel(tensor_model_parallel_size=world_size) + group = get_tensor_model_parallel_group().device_group + cpu_group = get_tensor_model_parallel_group().cpu_group + pynccl_comm = get_tensor_model_parallel_group().pynccl_comm + pymscclpp_comm = get_tensor_model_parallel_group().pymscclpp_comm + dist.barrier() + profile = False + dtype = torch.bfloat16 + ctx = get_torch_prof_ctx(profile) + result = [] + + with ctx: + for i in range(10, 20): + sz = 2**i + if sz * dtype.itemsize > 2**20: + break + inp_randn = torch.randint(1, 16, (sz,), dtype=dtype, device=device) + + memory = torch.empty_like(inp_randn) + memory_out = torch.empty_like(memory) + torch_eager_output, torch_eager_time = _bench_eager_time( + lambda inp: torch_allreduce(inp, group), inp_randn + ) + msccl_eager_output, msccl_eager_time = _bench_eager_time( + lambda inp: msccl_allreduce(inp, pymscclpp_comm), inp_randn + ) + msccl_graph_output, msccl_graph_time = _bench_graph_time( + lambda inp: msccl_allreduce(inp, pymscclpp_comm), inp_randn + ) + # since pynccl is inplace op, this return result is not correct if graph loop > 1 + _, pynccl_graph_time = _bench_graph_time( + lambda inp: pynccl_allreduce(inp, pynccl_comm), inp_randn + ) + torch.testing.assert_close(torch_eager_output, msccl_graph_output) + torch.testing.assert_close(torch_eager_output, msccl_eager_output) + result.append( + { + "msg_size": human_readable_size(inp_randn.nbytes), + "torch eager time": torch_eager_time, + "msccl eager time": msccl_eager_time, + "msccl graph time": msccl_graph_time, + "pynccl graph time": pynccl_graph_time, + } + ) + if rank == 0: + print(f"sz={sz}, dtype={dtype}: correctness check PASS!") + if rank == 0: + print_markdown_table(result) + if profile: + prof_dir = f"prof/msccl" + os.makedirs(prof_dir, exist_ok=True) + ctx.export_chrome_trace(f"{prof_dir}/trace_rank{dist.get_rank()}.json.gz") diff --git a/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py b/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py new file mode 100644 index 000000000..b61240b05 --- /dev/null +++ b/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py @@ -0,0 +1,403 @@ +import itertools +import math + +import cudnn +import torch +import torch.utils.benchmark as benchmark +from flashinfer import BatchDecodeWithPagedKVCacheWrapper + +from sglang.srt.layers.attention.triton_ops.decode_attention import decode_attention_fwd +from sglang.srt.utils import should_use_tensor_core + + +def benchmark_forward( + fn, + *inputs, + repeats=10, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + def amp_wrapper(*inputs, **kwinputs): + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + fn(*inputs, **kwinputs) + + t = benchmark.Timer( + stmt="fn_amp(*inputs, **kwinputs)", + globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + return t, m + + +def time_fwd(func, *args, **kwargs): + time_f = benchmark_forward(func, *args, **kwargs) + return time_f[1].mean * 1e6 + + +def decode_attention_sglang( + q, + kv_data, + batch_size, + kv_len, + head_num_q, + head_num_kv, + head_dim, + num_kv_splits, + warmup=10, +): + + k_buffer = kv_data[0].view(-1, head_num_kv, head_dim) + v_buffer = kv_data[1].view(-1, head_num_kv, head_dim) + o = torch.empty_like(q) + total_tokens = batch_size * kv_len + req_to_token = torch.arange(0, total_tokens).to(0).int().view(batch_size, kv_len) + b_req_idx = torch.arange(0, batch_size).to(0).int() + b_seq_len = torch.full((batch_size,), kv_len, dtype=torch.int32, device="cuda") + max_len_in_batch = kv_len + sm_scale = 1.0 / (head_dim**0.5) + + attn_logits = torch.empty( + (batch_size, head_num_q, num_kv_splits, head_dim + 1), + dtype=torch.float32, + device="cuda", + ) + + for _ in range(warmup): + decode_attention_fwd( + q, + k_buffer, + v_buffer, + o, + req_to_token, + b_req_idx, + b_seq_len, + attn_logits, + num_kv_splits, + sm_scale, + ) + + f = time_fwd( + decode_attention_fwd, + q, + k_buffer, + v_buffer, + o, + req_to_token, + b_req_idx, + b_seq_len, + attn_logits, + num_kv_splits, + sm_scale, + ) + + return f, o + + +def decode_attention_flashinfer(dtype, head_num_q, head_num_kv): + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda") + use_tensor_cores = should_use_tensor_core( + kv_cache_dtype=dtype, + num_attention_heads=head_num_q, + num_kv_heads=head_num_kv, + ) + flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, "NHD", use_tensor_cores=use_tensor_cores + ) + + class FlashinferAttention(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + kv_data, + batch_size, + kv_len, + head_num_q, + head_num_kv, + head_dim, + dtype, + warmup=10, + ): + total_tokens = batch_size * kv_len + kv_indptr = torch.arange(0, batch_size + 1).to(0).int() * kv_len + kv_indices = torch.arange(0, total_tokens).to(0).int() + kv_last_page_len = torch.full( + (batch_size,), 1, dtype=torch.int32, device="cuda" + ) + + flashinfer_decode_wrapper.end_forward() + flashinfer_decode_wrapper.begin_forward( + kv_indptr, + kv_indices, + kv_last_page_len, + head_num_q, + head_num_kv, + head_dim, + 1, + pos_encoding_mode="NONE", + data_type=dtype, + ) + + for _ in range(warmup): + o = flashinfer_decode_wrapper.forward( + q.contiguous().view(-1, head_num_q, head_dim), kv_data + ) + + f = time_fwd( + flashinfer_decode_wrapper.forward, + q.contiguous().view(-1, head_num_q, head_dim), + kv_data, + ) + + return f, o + + return FlashinferAttention + + +def convert_to_cudnn_type(torch_type): + if torch_type == torch.float16: + return cudnn.data_type.HALF + elif torch_type == torch.bfloat16: + return cudnn.data_type.BFLOAT16 + elif torch_type == torch.float32: + return cudnn.data_type.FLOAT + elif torch_type == torch.int32: + return cudnn.data_type.INT32 + elif torch_type == torch.int64: + return cudnn.data_type.INT64 + else: + raise ValueError("Unsupported tensor data type.") + + +def decode_attention_cudnn( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype, warmup=10 +): + # Prepare data: continuous q,k,v + dims_q = (batch_size, head_num_q, 1, head_dim) + strides_q = (head_num_q * head_dim, head_dim, head_num_q * head_dim, 1) + q_gpu = q.as_strided(dims_q, strides_q) + o_gpu = ( + torch.empty(batch_size * head_num_q * head_dim) + .half() + .cuda() + .as_strided(dims_q, strides_q) + ) + + dims_kv = (batch_size, head_num_kv, kv_len, head_dim) + strides_kv = ( + kv_len * head_num_kv * head_dim, + head_dim, + head_num_kv * head_dim, + 1, + ) + k_gpu = kv_data[0].as_strided(dims_kv, strides_kv) + v_gpu = kv_data[1].as_strided(dims_kv, strides_kv) + + seq_len_q_gpu = torch.full((batch_size, 1, 1, 1), 1, device="cuda") + seq_len_kv_gpu = torch.full((batch_size, 1, 1, 1), kv_len, device="cuda") + attn_scale = 1.0 / (head_dim**0.5) + + # Prepare data: paged k,v + block_size = 1 + blocks_per_batch = math.ceil(kv_len / block_size) + # [num_blocks, head_num_kv, block_size, head_dim], num_blocks = batch_size * blocks_per_batch + container_k_gpu = torch.cat(k_gpu.chunk(blocks_per_batch, dim=2), dim=0) + container_v_gpu = torch.cat(v_gpu.chunk(blocks_per_batch, dim=2), dim=0) + page_table_k_gpu = ( + torch.linspace( + 0, + batch_size * blocks_per_batch - 1, + batch_size * blocks_per_batch, + device="cuda", + dtype=torch.int32, + ) + .reshape(blocks_per_batch, 1, batch_size, 1) + .transpose(0, 2) + ) + page_table_v_gpu = page_table_k_gpu.clone() + + graph = cudnn.pygraph( + io_data_type=convert_to_cudnn_type(dtype), + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + q = graph.tensor_like(q_gpu) + container_k = graph.tensor_like(container_k_gpu) + container_v = graph.tensor_like(container_v_gpu) + page_table_k = graph.tensor_like(page_table_k_gpu) + page_table_v = graph.tensor_like(page_table_v_gpu) + + seq_len_q = graph.tensor_like(seq_len_q_gpu) + seq_len_kv = graph.tensor_like(seq_len_kv_gpu) + + o, _ = graph.sdpa( + name="sdpa", + q=q, + k=container_k, # Container K: non contiguous container with K blocks + v=container_v, # Container V: non contiguous container with V blocks + is_inference=True, + attn_scale=attn_scale, + use_causal_mask=False, + use_padding_mask=True, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + paged_attention_k_table=page_table_k, # Page Table K: Tensor containing offsets to the container with K blocks + paged_attention_v_table=page_table_v, # Page Table V: Tensor containing offsets to the container with V blocks + paged_attention_max_seq_len_kv=kv_len, # The maximum sequence length for K caches (this is optional, but recommended) + ) + + o.set_output(True).set_dim(dims_q).set_stride(strides_q) + + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A]) + graph.check_support() + graph.build_plans() + + workspace = torch.empty( + graph.get_workspace_size(), device="cuda", dtype=torch.uint8 + ) + + variant_pack = { + q: q_gpu, + container_k: container_k_gpu, + container_v: container_v_gpu, + page_table_k: page_table_k_gpu, + page_table_v: page_table_v_gpu, + seq_len_q: seq_len_q_gpu, + seq_len_kv: seq_len_kv_gpu, + o: o_gpu, + } + + for _ in range(warmup): + graph.execute(variant_pack, workspace) + + f = time_fwd( + graph.execute, + variant_pack, + workspace, + ) + + return f, o_gpu.squeeze(dim=2) + + +def calculate_diff(): + + dtype = torch.float16 + batch_size = 64 + kv_len = 4096 + head_num_q = 64 + head_num_kv = 8 + head_dim = 128 + + q = torch.randn(batch_size, head_num_q, head_dim, dtype=dtype, device="cuda") + kv_data = ( + torch.randn( + batch_size * kv_len, head_num_kv, head_dim, dtype=dtype, device="cuda" + ), + torch.randn( + batch_size * kv_len, head_num_kv, head_dim, dtype=dtype, device="cuda" + ), + ) + + _, output_sglang = decode_attention_sglang( + q, + kv_data, + batch_size, + kv_len, + head_num_q, + head_num_kv, + head_dim, + num_kv_splits=8, + ) + + attn_flashinfer = decode_attention_flashinfer(dtype, head_num_q, head_num_kv).apply + _, output_flashinfer = attn_flashinfer( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype + ) + + _, output_cudnn = decode_attention_cudnn( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype + ) + + print(f"SGLang output={output_sglang}") + print(f"FlashInfer output={output_flashinfer}") + print(f"cuDNN output={output_cudnn}") + if torch.allclose(output_sglang, output_flashinfer, atol=1e-2, rtol=1e-2): + print("✅ SGLang[Triton] and FlashInfer match") + else: + print("❌ SGLang[Triton] and FlashInfer differ") + + if torch.allclose(output_sglang, output_cudnn, atol=1e-2, rtol=1e-2): + print("✅ SGLang[Triton] and cuDNN match") + else: + print("❌ SGLang[Triton] and cuDNN differ") + + +if __name__ == "__main__": + calculate_diff() + + head_dim = 128 + dtype = torch.float16 + batch_size_range = [2**i for i in range(0, 8, 2)] + kv_len_range = [2**i for i in range(6, 13, 1)] + configs = list(itertools.product(batch_size_range, kv_len_range)) + + for head_num_q, head_num_kv in [[32, 32], [64, 8], [40, 8]]: + attn_flashinfer = decode_attention_flashinfer( + dtype, head_num_q, head_num_kv + ).apply + for batch_size, kv_len in configs: + q = torch.randn( + batch_size, head_num_q, head_dim, dtype=dtype, device="cuda" + ) + kv_data = ( + torch.randn( + batch_size * kv_len, + head_num_kv, + head_dim, + dtype=dtype, + device="cuda", + ), + torch.randn( + batch_size * kv_len, + head_num_kv, + head_dim, + dtype=dtype, + device="cuda", + ), + ) + us_cudnn, output_cudnn = decode_attention_cudnn( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype + ) + us_sglang, output_sglang = decode_attention_sglang( + q, + kv_data, + batch_size, + kv_len, + head_num_q, + head_num_kv, + head_dim, + num_kv_splits=8, + ) + us_flashinfer, _ = attn_flashinfer( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype + ) + print( + head_num_q, + " ", + head_num_kv, + " ", + batch_size, + " ", + kv_len, + " ", + us_cudnn, + " ", + us_sglang, + " ", + us_flashinfer, + ) diff --git a/benchmark/kernels/deepep/deepep_utils.py b/benchmark/kernels/deepep/deepep_utils.py new file mode 100644 index 000000000..169529ef6 --- /dev/null +++ b/benchmark/kernels/deepep/deepep_utils.py @@ -0,0 +1,218 @@ +# ADAPTED FROM https://github.com/deepseek-ai/DeepEP/blob/main/tests/utils.py + +import os +import sys +from typing import Optional + +import numpy as np +import torch +import torch.distributed as dist + + +def init_dist(local_rank: int, num_local_ranks: int, args): + ip = args.master_addr + port = args.master_port + num_nodes = args.nnodes + node_rank = args.node_rank + assert (num_local_ranks < 8 and num_nodes == 1) or num_local_ranks == 8 + + dist.init_process_group( + backend="nccl", + init_method=f"tcp://{ip}:{port}", + world_size=num_nodes * num_local_ranks, + rank=node_rank * num_local_ranks + local_rank, + ) + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.cuda.set_device(local_rank) + + return ( + dist.get_rank(), + dist.get_world_size(), + dist.new_group(list(range(num_local_ranks * num_nodes))), + ) + + +def calc_diff(x: torch.Tensor, y: torch.Tensor): + x, y = x.double() + 1, y.double() + 1 + denominator = (x * x + y * y).sum() + sim = 2 * (x * y).sum() / denominator + return (1 - sim).item() + + +def per_token_cast_to_fp8(x: torch.Tensor): + assert x.dim() == 2 and x.size(1) % 128 == 0 + m, n = x.shape + x_view = x.view(m, -1, 128) + x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) + return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( + m, n + ), (x_amax / 448.0).view(m, -1) + + +def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor): + x_fp32 = x_fp8.to(torch.float32).view(x_fp8.size(0), -1, 128) + x_scales = x_scales.view(x_fp8.size(0), -1, 1) + return (x_fp32 * x_scales).view(x_fp8.shape).to(torch.bfloat16) + + +def inplace_unique(x: torch.Tensor, num_slots: int): + assert x.dim() == 2 + mask = x < 0 + x_padded = x.masked_fill(mask, num_slots) + bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device) + bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded)) + bin_count = bin_count[:, :num_slots] + sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True) + sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1) + sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values + x[:, :].fill_(-1) + valid_len = min(num_slots, x.size(1)) + x[:, :valid_len] = sorted_bin_idx[:, :valid_len] + + +def create_grouped_scores( + scores: torch.Tensor, group_idx: torch.Tensor, num_groups: int +): + num_tokens, num_experts = scores.shape + scores = scores.view(num_tokens, num_groups, -1) + mask = torch.zeros((num_tokens, num_groups), dtype=torch.bool, device=scores.device) + mask = mask.scatter_(1, group_idx, True).unsqueeze(-1).expand_as(scores) + return (scores * mask).view(num_tokens, num_experts) + + +def bench(fn, num_warmups: int = 20, num_tests: int = 30, post_fn=None): + # Flush L2 cache with 256 MB data + torch.cuda.synchronize() + cache = torch.empty(int(256e6 // 4), dtype=torch.int, device="cuda") + + # Warmup + for _ in range(num_warmups): + fn() + + # Flush L2 + cache.zero_() + + # Testing + start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] + for i in range(num_tests): + # Record + start_events[i].record() + fn() + end_events[i].record() + if post_fn is not None: + post_fn() + torch.cuda.synchronize() + + times = np.array( + [s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)] + )[1:] + return np.average(times), np.min(times), np.max(times) + + +class empty_suppress: + def __enter__(self): + return self + + def __exit__(self, *_): + pass + + +class suppress_stdout_stderr: + def __enter__(self): + self.outnull_file = open(os.devnull, "w") + self.errnull_file = open(os.devnull, "w") + + self.old_stdout_fileno_undup = sys.stdout.fileno() + self.old_stderr_fileno_undup = sys.stderr.fileno() + + self.old_stdout_fileno = os.dup(sys.stdout.fileno()) + self.old_stderr_fileno = os.dup(sys.stderr.fileno()) + + self.old_stdout = sys.stdout + self.old_stderr = sys.stderr + + os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup) + os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup) + + sys.stdout = self.outnull_file + sys.stderr = self.errnull_file + return self + + def __exit__(self, *_): + sys.stdout = self.old_stdout + sys.stderr = self.old_stderr + + os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup) + os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup) + + os.close(self.old_stdout_fileno) + os.close(self.old_stderr_fileno) + + self.outnull_file.close() + self.errnull_file.close() + + +def bench_kineto( + fn, + kernel_names, + num_tests: int = 30, + suppress_kineto_output: bool = False, + trace_path: Optional[str] = None, + barrier_comm_profiling: bool = False, +): + # Profile + suppress = suppress_stdout_stderr if suppress_kineto_output else empty_suppress + with suppress(): + schedule = torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1) + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule + ) as prof: + for i in range(2): + # NOTES: use a large kernel and a barrier to eliminate the unbalanced CPU launch overhead + if barrier_comm_profiling: + lhs = torch.randn((8192, 8192), dtype=torch.float, device="cuda") + rhs = torch.randn((8192, 8192), dtype=torch.float, device="cuda") + lhs @ rhs + dist.all_reduce(torch.ones(1, dtype=torch.float, device="cuda")) + for _ in range(num_tests): + fn() + prof.step() + + # Parse the profiling table + assert isinstance(kernel_names, str) or isinstance(kernel_names, tuple) + is_tupled = isinstance(kernel_names, tuple) + prof_lines = ( + prof.key_averages() + .table(sort_by="cuda_time_total", max_name_column_width=100) + .split("\n") + ) + kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names + assert all([isinstance(name, str) for name in kernel_names]) + for name in kernel_names: + assert ( + sum([name in line for line in prof_lines]) == 1 + ), f"Errors of the kernel {name} in the profiling table" + + # Save chrome traces + if trace_path is not None: + prof.export_chrome_trace(trace_path) + + # Return average kernel times + units = {"ms": 1e3, "us": 1e6} + kernel_times = [] + for name in kernel_names: + for line in prof_lines: + if name in line: + time_str = line.split()[-2] + for unit, scale in units.items(): + if unit in time_str: + kernel_times.append(float(time_str.replace(unit, "")) / scale) + break + break + return tuple(kernel_times) if is_tupled else kernel_times[0] + + +def hash_tensor(t: torch.Tensor): + return t.view(torch.int64).sum().item() diff --git a/benchmark/kernels/deepep/tuning_deepep.py b/benchmark/kernels/deepep/tuning_deepep.py new file mode 100644 index 000000000..bb900a875 --- /dev/null +++ b/benchmark/kernels/deepep/tuning_deepep.py @@ -0,0 +1,476 @@ +# MODIFIED FROM https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_internode.py + +""" +Example usage: +python tuning_deepep.py --nnodes 4 --node-rank $MY_NODE_RANK --master-addr 1.2.3.4 +Then check `deepep_tuned.json` +""" + +import argparse +import json +import time +from copy import deepcopy +from pathlib import Path + +# noinspection PyUnresolvedReferences +import deep_ep +import torch +import torch.distributed as dist +from deepep_utils import ( + bench, + calc_diff, + create_grouped_scores, + init_dist, + inplace_unique, + per_token_cast_back, + per_token_cast_to_fp8, +) + + +def test_main( + num_sms: int, + local_rank: int, + num_local_ranks: int, + num_ranks: int, + num_nodes: int, + rank: int, + buffer: deep_ep.Buffer, + group: dist.ProcessGroup, + args, +): + # Settings + num_tokens, hidden, num_topk_groups, num_topk, num_experts = ( + 4096, + 7168, + min(num_nodes, 4), + 8, + (256 // num_ranks) * num_ranks, + ) + assert num_experts % num_ranks == 0 and num_local_ranks == 8 + if local_rank == 0: + print( + f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk_groups={num_topk_groups}, num_topk={num_topk}", + flush=True, + ) + + # Random data + x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank + x_pure_rand = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + x_e4m3 = per_token_cast_to_fp8(x) + scores = ( + torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs() + + 1 + ) + group_scores = scores.view(num_tokens, num_nodes, -1).amax(dim=-1) + group_idx = torch.topk( + group_scores, k=num_topk_groups, dim=-1, sorted=False + ).indices + masked_scores = create_grouped_scores(scores, group_idx, num_nodes) + topk_idx = torch.topk(masked_scores, num_topk, dim=-1, largest=True, sorted=False)[ + 1 + ] + topk_weights = ( + torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") * rank + ) + topk_weights_pure_rand = torch.randn( + (num_tokens, num_topk), dtype=torch.float32, device="cuda" + ) + rank_idx = topk_idx // (num_experts // num_ranks) + rank_idx.masked_fill_(topk_idx == -1, -1) + inplace_unique(rank_idx, num_ranks) + rdma_rank_idx = rank_idx // num_local_ranks + rdma_rank_idx.masked_fill_(rank_idx == -1, -1) + inplace_unique(rdma_rank_idx, num_nodes) + + # RDMA dispatch counts + rdma_idx = topk_idx // (num_experts // num_nodes) + rdma_idx.masked_fill_(topk_idx == -1, -1) + inplace_unique(rdma_idx, num_nodes) + num_rdma_token_sent = rdma_idx.ne(-1).sum().item() + + # Expert meta + num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda") + for i in range(num_experts): + num_tokens_per_expert[i] = (topk_idx == i).sum() + gbl_num_tokens_per_expert = num_tokens_per_expert.clone() + dist.all_reduce(gbl_num_tokens_per_expert, group=group) + + # Rank layout meta + num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda") + num_tokens_per_rdma_rank = torch.empty((num_nodes,), dtype=torch.int, device="cuda") + token_idx_in_rank = torch.full( + (num_ranks, num_tokens), -1, dtype=torch.long, device="cuda" + ) + for i in range(num_ranks): + num_tokens_per_rank[i] = (rank_idx == i).sum() + token_sel = (rank_idx == i).max(dim=-1)[0] + count = token_sel.sum().item() + tokens = torch.sort(token_sel.to(torch.int), descending=True)[1] + tokens[:count] = torch.sort(tokens[:count])[0] + token_idx_in_rank[i][tokens[:count]] = torch.arange( + count, dtype=torch.long, device="cuda" + ) + for i in range(num_nodes): + num_tokens_per_rdma_rank[i] = (rdma_rank_idx == i).sum() + token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int) + is_token_in_rank = token_idx_in_rank >= 0 + gbl_num_tokens_per_rank = num_tokens_per_rank.clone() + dist.all_reduce(gbl_num_tokens_per_rank, group=group) + + ( + ref_num_tokens_per_rank, + ref_num_tokens_per_rdma_rank, + ref_num_tokens_per_expert, + ref_is_token_in_rank, + _, + ) = buffer.get_dispatch_layout(topk_idx, num_experts) + assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank) + assert torch.allclose(ref_num_tokens_per_rdma_rank, num_tokens_per_rdma_rank) + assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert) + assert torch.allclose(ref_is_token_in_rank, is_token_in_rank) + t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0] + if local_rank == 0: + print(f"[layout] Kernel performance: {t * 1000:.3f} ms", flush=True) + print("", flush=True) + group.barrier() + time.sleep(1) + + # Config + rdma_buffer_size, nvl_buffer_size = 128, (720 if num_ranks in (144, 160) else 512) + config = deep_ep.Config(num_sms, 8, nvl_buffer_size, 16, rdma_buffer_size) + + # Test dispatch + # noinspection PyShadowingNames + def check_data(check_x, recv_gbl_rank_prefix_sum): + assert torch.allclose(check_x.amin(dim=1), check_x.amax(dim=1)) + check_start = 0 + for i in range(num_ranks): + check_end = recv_gbl_rank_prefix_sum[i].item() + assert (check_x[check_start:check_end, :].int() - i).sum().item() == 0 + check_start = check_end + + for previous_mode in (False, True): + for async_mode in (False, True): + for current_x in (x_pure_rand, x, x_e4m3): + for with_topk in (False, True): + if local_rank == 0: + print( + f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...', + flush=True, + end="", + ) + dispatch_args = { + "x": current_x, + "num_tokens_per_rank": num_tokens_per_rank, + "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, + "is_token_in_rank": is_token_in_rank, + "num_tokens_per_expert": num_tokens_per_expert, + "config": config, + "async_finish": async_mode, + } + if with_topk: + dispatch_args.update( + { + "topk_idx": topk_idx, + "topk_weights": ( + topk_weights_pure_rand + if current_x is x_pure_rand + else topk_weights + ), + } + ) + if previous_mode: + dispatch_args.update({"previous_event": buffer.capture()}) + ( + recv_x, + recv_topk_idx, + recv_topk_weights, + recv_num_tokens_per_expert_list, + handle, + event, + ) = buffer.dispatch(**dispatch_args) + event.current_stream_wait() if async_mode else () + recv_x = ( + per_token_cast_back(*recv_x) + if isinstance(recv_x, tuple) + else recv_x + ) + + # Checks + recv_gbl_rank_prefix_sum = handle[-4] + assert gbl_num_tokens_per_rank[rank].item() == recv_x.size( + 0 + ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" + assert ( + gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist() + == recv_num_tokens_per_expert_list + ) + if current_x is not x_pure_rand: + check_data(recv_x, recv_gbl_rank_prefix_sum) + if with_topk: + # Check `topk_idx` + assert ( + recv_topk_idx.eq(-1) + | ( + (recv_topk_idx >= 0) + & (recv_topk_idx < (num_experts // num_ranks)) + ) + ).sum().item() == recv_topk_idx.numel() + for i, count in enumerate(recv_num_tokens_per_expert_list): + assert recv_topk_idx.eq(i).sum().item() == count + + # Check `topk_weights` + if current_x is not x_pure_rand: + recv_topk_weights[recv_topk_idx.eq(-1)] = ( + recv_topk_weights.amax(dim=1, keepdim=True).expand_as( + recv_topk_weights + )[recv_topk_idx.eq(-1)] + ) + check_data(recv_topk_weights, recv_gbl_rank_prefix_sum) + + # Test cached dispatch (must without top-k staffs) + if not with_topk: + dispatch_args = { + "x": current_x, + "handle": handle, + "config": config, + "async_finish": async_mode, + } + if previous_mode: + dispatch_args.update({"previous_event": buffer.capture()}) + recv_x, _, _, _, _, event = buffer.dispatch(**dispatch_args) + event.current_stream_wait() if async_mode else () + recv_x = ( + per_token_cast_back(*recv_x) + if isinstance(recv_x, tuple) + else recv_x + ) + if current_x is not x_pure_rand: + check_data(recv_x, recv_gbl_rank_prefix_sum) + + # Test combine + combine_args = { + "x": recv_x, + "handle": handle, + "config": config, + "async_finish": async_mode, + } + if with_topk: + combine_args.update({"topk_weights": recv_topk_weights}) + if previous_mode: + combine_args.update({"previous_event": buffer.capture()}) + combined_x, combined_topk_weights, event = buffer.combine( + **combine_args + ) + event.current_stream_wait() if async_mode else () + check_x = combined_x.float() / is_token_in_rank.sum( + dim=1 + ).unsqueeze(1) + ref_x = x_pure_rand if current_x is x_pure_rand else x + assert calc_diff(check_x, ref_x) < 5e-6 + if with_topk: + check_topk_weights = ( + combined_topk_weights + if (current_x is x_pure_rand) + else ( + combined_topk_weights + / is_token_in_rank.sum(dim=1).unsqueeze(1) + ) + ) + ref_topk_weights = ( + topk_weights_pure_rand + if current_x is x_pure_rand + else topk_weights + ) + assert calc_diff(check_topk_weights, ref_topk_weights) < 1e-9 + + # For later tuning + dispatch_bf16_rdma_send_bytes = num_rdma_token_sent * hidden * 2 + dispatch_bf16_nvl_recv_bytes = recv_x.numel() * 2 + combine_bf16_nvl_send_bytes = dispatch_bf16_nvl_recv_bytes + combine_bf16_rdma_recv_bytes = dispatch_bf16_rdma_send_bytes + + if local_rank == 0: + print(" passed", flush=True) + if local_rank == 0: + print("", flush=True) + + output_data = {} + + # Tune dispatch performance + best_dispatch_results = None + fp8_factor = (1 + 4 / 128) / 2 + for current_x in (x_e4m3, x): + best_time, best_results = 1e10, None + rdma_send_bytes = ( + (dispatch_bf16_rdma_send_bytes * fp8_factor) + if isinstance(current_x, tuple) + else dispatch_bf16_rdma_send_bytes + ) + nvl_recv_bytes = ( + (dispatch_bf16_nvl_recv_bytes * fp8_factor) + if isinstance(current_x, tuple) + else dispatch_bf16_nvl_recv_bytes + ) + for nvl_chunk_size in range(4, 33, 4): + for rdma_chunk_size in range(4, 33, 4): + config_kwargs = { + "num_sms": num_sms, + "num_max_nvl_chunked_send_tokens": nvl_chunk_size, + "num_max_nvl_chunked_recv_tokens": nvl_buffer_size, + "num_max_rdma_chunked_send_tokens": rdma_chunk_size, + "num_max_rdma_chunked_recv_tokens": rdma_buffer_size, + } + config = deep_ep.Config(**config_kwargs) + tune_args = {"x": current_x, "handle": handle, "config": config} + t = bench(lambda: buffer.dispatch(**tune_args))[0] + if t < best_time: + best_time, best_results = t, ( + num_sms, + nvl_chunk_size, + rdma_chunk_size, + config_kwargs, + ) + if local_rank == 0: + print( + f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {rdma_send_bytes / 1e9 / t:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ", + flush=True, + ) + if local_rank == 0: + print( + f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)', + flush=True, + ) + print("", flush=True) + is_fp8 = isinstance(current_x, tuple) + if is_fp8: + output_data["normal_dispatch"] = deepcopy(best_results[3]) + + if isinstance(current_x, tuple): + # Gather FP8 the best config from rank 0 + best_dispatch_results = torch.tensor( + [best_results[0], best_results[1], best_results[2]], + dtype=torch.int32, + device="cuda", + ) + all_best_fp8_results_list = [ + torch.zeros_like(best_dispatch_results) + for _ in range(torch.distributed.get_world_size()) + ] + dist.all_gather( + all_best_fp8_results_list, best_dispatch_results, group=group + ) + best_dispatch_results = all_best_fp8_results_list[0].tolist() + dispatch_config = deep_ep.Config( + best_dispatch_results[0], + best_dispatch_results[1], + nvl_buffer_size, + best_dispatch_results[2], + rdma_buffer_size, + ) + + dispatch_args = { + "x": x, + "num_tokens_per_rank": num_tokens_per_rank, + "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, + "is_token_in_rank": is_token_in_rank, + "num_tokens_per_expert": num_tokens_per_expert, + "config": dispatch_config if dispatch_config is not None else config, + } + recv_x, _, _, _, handle, _ = buffer.dispatch(**dispatch_args) + + # Tune combine performance + best_time, best_results = 1e10, None + for nvl_chunk_size in range(1, 5, 1): + for rdma_chunk_size in range(8, 33, 4): + config_kwargs = { + "num_sms": num_sms, + "num_max_nvl_chunked_send_tokens": nvl_chunk_size, + "num_max_nvl_chunked_recv_tokens": nvl_buffer_size, + "num_max_rdma_chunked_send_tokens": rdma_chunk_size, + "num_max_rdma_chunked_recv_tokens": rdma_buffer_size, + } + config = deep_ep.Config(**config_kwargs) + tune_args = {"x": recv_x, "handle": handle, "config": config} + t = bench(lambda: buffer.combine(**tune_args))[0] + if local_rank == 0: + print( + f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {combine_bf16_rdma_recv_bytes / 1e9 / t:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ", + flush=True, + ) + if t < best_time: + best_time, best_results = t, ( + num_sms, + nvl_chunk_size, + rdma_chunk_size, + config_kwargs, + ) + + if local_rank == 0: + print( + f"[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {combine_bf16_rdma_recv_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)", + flush=True, + ) + print("", flush=True) + output_data["normal_combine"] = deepcopy(best_results[3]) + + if rank == 0 and local_rank == 0: + _write_output(args, output_data) + + +def _write_output(args, output_data): + text = json.dumps(output_data, indent=4) + output_path = args.output_path + print(f"Write to {output_path} with {text}") + Path(output_path).write_text(text) + + +# noinspection PyUnboundLocalVariable +def test_loop(local_rank: int, num_local_ranks: int, args): + num_nodes = args.nnodes + rank, num_ranks, group = init_dist(local_rank, num_local_ranks, args) + + num_sms = args.num_sms + num_qps_per_rank = num_sms // 2 + + buffer = deep_ep.Buffer( + group, + int(1e9), + int(1e9), + low_latency_mode=False, + num_qps_per_rank=num_qps_per_rank, + ) + assert num_local_ranks == 8 and num_ranks > 8 + torch.manual_seed(rank) + + for i in (num_sms,): + test_main( + i, + local_rank, + num_local_ranks, + num_ranks, + num_nodes, + rank, + buffer, + group, + args, + ) + if local_rank == 0: + print("", flush=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-sms", type=int, default=24) + parser.add_argument("--output-path", type=str, default="deepep_tuned.json") + parser.add_argument("--nnodes", type=int, default=1) + parser.add_argument("--node-rank", type=int, default=0) + parser.add_argument("--master-addr", type=str, default="127.0.0.1") + parser.add_argument("--master-port", type=int, default=8361) + args = parser.parse_args() + print(f"Start system with {args=}") + + num_processes = 8 + torch.multiprocessing.spawn( + test_loop, args=(num_processes, args), nprocs=num_processes + ) diff --git a/benchmark/kernels/deepseek/README.md b/benchmark/kernels/deepseek/README.md new file mode 100644 index 000000000..fba775f23 --- /dev/null +++ b/benchmark/kernels/deepseek/README.md @@ -0,0 +1,19 @@ +## DeepSeek kernels benchmark + + +### Prerequisites +- You should install [DeepGemm](https://github.com/deepseek-ai/DeepGEMM) from source before run `benchmark_deepgemm_fp8_gemm.py` and `benchmark_deepgemm_fp8_group_gemm.py`. + +### Benchmark +- `benchmark_deepgemm_fp8_gemm.py` + ```bash + python benchmark_deepgemm_fp8_gemm.py --run_correctness --tp_size 1 + ``` + +- `benchmark_deepgemm_fp8_group_gemm.py` + ```bash + python benchmark_deepgemm_fp8_group_gemm.py --run_correctness --tp_size 1 + ``` + + - You can use the `--run_correctness` parameter to verify all kernels results's correctness. + - You can use the `--tp_size` parameter to benchmark all FP8 w8a8 block-wise matrix multiplications involved in DeepSeek V3/R1 under the current tensor parallelism (TP) setting. This benchmark compares DeepSeek's open-source [DeepGemm](https://github.com/deepseek-ai/DeepGEMM) implementation with SGLang's and VLLM Triton implementation. diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py new file mode 100644 index 000000000..f93732154 --- /dev/null +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py @@ -0,0 +1,400 @@ +from typing import Tuple + +import deep_gemm +import tilelang +import tilelang.language as T +import torch +import triton +from deep_gemm import ceil_div, get_col_major_tma_aligned_tensor +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + w8a8_block_fp8_matmul as vllm_w8a8_block_fp8_matmul, +) + +from sglang.srt.layers.quantization.fp8_kernel import ( + w8a8_block_fp8_matmul_deepgemm as w8a8_block_fp8_matmul, +) + + +# Adapted from https://github.com/tile-ai/tilelang/blob/a8cfdce92795cb861c9033573534653ee040b5ed/examples/deepseek_deepgemm/example_deepgemm_fp8_2xAcc.py#L1 +def tl_gemm( + M, + N, + K, + in_dtype, + out_dtype, + accum_dtype, +): + assert in_dtype in [ + "e4m3_float8", + ], "Currently only e4m3_float8 is supported" + assert out_dtype in [ + "bfloat16", + "float16", + ], "Currently only bfloat16 and float16 are supported" + + TILE_SIZE = (128, 128, 128) + block_M = TILE_SIZE[0] + block_N = TILE_SIZE[1] + block_K = TILE_SIZE[2] + + A_shape = (M, K) + Scales_A_shape = (M, T.ceildiv(K, block_K)) + B_shape = (N, K) + Scales_B_shape = (T.ceildiv(N, block_N), T.ceildiv(K, block_K)) + A_shared_shape = (block_M, block_K) + B_shared_shape = (block_N, block_K) + C_shared_shape = (block_M, block_N) + + @T.prim_func + def main( + A: T.Buffer(A_shape, in_dtype), + scales_a: T.Buffer(Scales_A_shape, "float32"), + B: T.Buffer(B_shape, in_dtype), + scales_b: T.Buffer(Scales_B_shape, "float32"), + C: T.Buffer((M, N), out_dtype), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as ( + bx, + by, + ): + + A_shared = T.alloc_shared(A_shared_shape, in_dtype) + B_shared = T.alloc_shared(B_shared_shape, in_dtype) + C_shared = T.alloc_shared(C_shared_shape, out_dtype) + Scale_C_shared = T.alloc_shared((block_M), "float32") + C_local = T.alloc_fragment(C_shared_shape, accum_dtype) + C_local_accum = T.alloc_fragment(C_shared_shape, accum_dtype) + + # Improve L2 Cache + T.use_swizzle(panel_size=10) + + T.clear(C_local) + T.clear(C_local_accum) + K_iters = T.ceildiv(K, block_K) + for k in T.Pipelined(K_iters, num_stages=4): + # Load A into shared memory + T.copy(A[by * block_M, k * block_K], A_shared) + # Load B into shared memory + T.copy(B[bx * block_N, k * block_K], B_shared) + # Load scale into shared memory + Scale_B = scales_b[bx, k] + for i in T.Parallel(block_M): + Scale_C_shared[i] = scales_a[by * block_M + i, k] * Scale_B + + T.gemm(A_shared, B_shared, C_local, transpose_B=True) + # Promote to enable 2xAcc + for i, j in T.Parallel(block_M, block_N): + C_local_accum[i, j] += C_local[i, j] * Scale_C_shared[i] + T.clear(C_local) + # TMA store + T.copy(C_local_accum, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return main + + +def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 and x.size(1) % 128 == 0 + m, n = x.shape + x_view = x.view(m, -1, 128) + x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) + return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( + m, n + ), (x_amax / 448.0).view(m, -1) + + +def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + m, n = x.shape + x_padded = torch.zeros( + (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device + ) + x_padded[:m, :n] = x + x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) + x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn) + return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view( + x_view.size(0), x_view.size(2) + ) + + +def fp8_gemm_deepgemm( + x_fp8: torch.Tensor, + x_scale: torch.Tensor, + y_fp8: torch.Tensor, + y_scale: torch.Tensor, + m: int, + n: int, + k: int, +): + """DeepGEMM implementation of FP8 GEMM""" + out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + + # Run DeepGEMM kernel + deep_gemm.gemm_fp8_fp8_bf16_nt((x_fp8, x_scale), (y_fp8, y_scale), out) + return out + + +def fp8_gemm_sglang( + x_fp8: torch.Tensor, + x_scale: torch.Tensor, + y_fp8: torch.Tensor, + y_scale: torch.Tensor, + m: int, + n: int, + k: int, +): + """SGLang implementation of FP8 GEMM""" + block_size = [128, 128] # Matches the block size in per_block_cast_to_fp8 + + # Run SGLang kernel + out = w8a8_block_fp8_matmul( + x_fp8, y_fp8, x_scale, y_scale, block_size, torch.bfloat16 + ) + return out + + +def fp8_gemm_vllm( + x_fp8: torch.Tensor, + x_scale: torch.Tensor, + y_fp8: torch.Tensor, + y_scale: torch.Tensor, + m: int, + n: int, + k: int, +): + """vLLM implementation of FP8 GEMM""" + block_size = [128, 128] # Matches the block size in per_block_cast_to_fp8 + + # Run vLLM kernel + out = vllm_w8a8_block_fp8_matmul( + x_fp8, y_fp8, x_scale, y_scale, block_size, torch.bfloat16 + ) + return out + + +def calculate_diff(m: int, n: int, k: int): + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + + x_fp8, x_scale = per_token_cast_to_fp8(x.clone()) + y_fp8, y_scale = per_block_cast_to_fp8(y.clone()) + x_scale_col_major = get_col_major_tma_aligned_tensor(x_scale.clone()) + + out_deepgemm = fp8_gemm_deepgemm( + x_fp8.clone(), + x_scale_col_major.clone(), + y_fp8.clone(), + y_scale.clone(), + m, + n, + k, + ) + out_sglang = fp8_gemm_sglang( + x_fp8.clone(), x_scale.clone(), y_fp8.clone(), y_scale.clone(), m, n, k + ) + + tilelang_func = tl_gemm(m, n, k, "e4m3_float8", "bfloat16", "float32") + tilelang_kernel = tilelang.compile(tilelang_func, out_idx=[-1]) + out_tilelang = tilelang_kernel( + x_fp8.clone(), x_scale.clone(), y_fp8.clone(), y_scale.clone() + ) + + diff_sglang_deepgemm = torch.abs(out_deepgemm - out_sglang).mean().item() + diff_tilelang_deepgemm = torch.abs(out_deepgemm - out_tilelang).mean().item() + diff_tilelang_sglang = torch.abs(out_tilelang - out_sglang).mean().item() + + print(f"Shape m={m}, n={n}, k={k}:") + print(f"DeepGEMM output: {out_deepgemm[0, 0:5]}") + print(f"SGLang output: {out_sglang[0, 0:5]}") + print(f"TileLang output: {out_tilelang[0, 0:5]}") + print(f"Mean absolute difference (SGLang-DeepGEMM): {diff_sglang_deepgemm}") + print(f"Mean absolute difference (TileLang-DeepGEMM): {diff_tilelang_deepgemm}") + print(f"Mean absolute difference (TileLang-SGLang): {diff_tilelang_sglang}") + + sglang_deepgemm_match = torch.allclose( + out_deepgemm, out_sglang, atol=1e-2, rtol=1e-2 + ) + tilelang_deepgemm_match = torch.allclose( + out_deepgemm, out_tilelang, atol=1e-2, rtol=1e-2 + ) + tilelang_sglang_match = torch.allclose( + out_tilelang, out_sglang, atol=1e-2, rtol=1e-2 + ) + + if sglang_deepgemm_match and tilelang_deepgemm_match and tilelang_sglang_match: + print("✅ All implementations match\n") + else: + print("❌ Some implementations differ:") + print(f" - SGLang vs DeepGEMM: {'✅' if sglang_deepgemm_match else '❌'}") + print(f" - TileLang vs DeepGEMM: {'✅' if tilelang_deepgemm_match else '❌'}") + print(f" - TileLang vs SGLang: {'✅' if tilelang_sglang_match else '❌'}\n") + + +def get_weight_shapes(tp_size): + # cannot TP + total = [ + (512 + 64, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (7168, 16384), + (7168, 18432), + ] + # N can TP + n_tp = [ + (18432 * 2, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (24576, 1536), + (4096, 7168), + ] + # K can TP + k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] + + weight_shapes = [] + for t in total: + weight_shapes.append(t) + for n_t in n_tp: + new_t = (n_t[0] // tp_size, n_t[1]) + weight_shapes.append(new_t) + for k_t in k_tp: + new_t = (k_t[0], k_t[1] // tp_size) + weight_shapes.append(new_t) + + return weight_shapes + + +def create_benchmark_configs(tp_size): + configs = [] + weight_shapes = get_weight_shapes(tp_size) + batch_sizes = [8, 16, 32, 64, 128, 256, 1024, 2048, 4096] + + for n, k in weight_shapes: + for m in batch_sizes: + configs.append((m, n, k, tp_size)) + + return configs + + +def get_benchmark(tp_size): + all_configs = create_benchmark_configs(tp_size) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["m", "n", "k", "tp_size"], + x_vals=[list(config) for config in all_configs], + line_arg="provider", + line_vals=["deepgemm", "sglang", "tilelang"], + line_names=["DeepGEMM", "SGLang", "TileLang"], + styles=[("blue", "-"), ("red", "-"), ("green", "-")], + ylabel="ms", + plot_name=f"fp8-gemm-performance-comparison-tp{tp_size}", + args={}, + ) + ) + def benchmark(m, n, k, tp_size, provider): + print(f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), Provider: {provider}") + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + + # Preprocess data before benchmarking + x_fp8, x_scale = per_token_cast_to_fp8(x) + y_fp8, y_scale = per_block_cast_to_fp8(y) + x_scale_col_major = get_col_major_tma_aligned_tensor(x_scale.clone()) + + quantiles = [0.5, 0.2, 0.8] + + if provider == "deepgemm": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: fp8_gemm_deepgemm( + x_fp8.clone(), + x_scale_col_major.clone(), + y_fp8.clone(), + y_scale.clone(), + m, + n, + k, + ), + quantiles=quantiles, + ) + elif provider == "sglang": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: fp8_gemm_sglang( + x_fp8.clone(), + x_scale.clone(), + y_fp8.clone(), + y_scale.clone(), + m, + n, + k, + ), + quantiles=quantiles, + ) + else: # tilelang + tilelang_func = tl_gemm(m, n, k, "e4m3_float8", "bfloat16", "float32") + tilelang_kernel = tilelang.compile(tilelang_func, out_idx=[-1]) + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: tilelang_kernel( + x_fp8.clone(), + x_scale.clone(), + y_fp8.clone(), + y_scale.clone(), + ), + quantiles=quantiles, + ) + + # Calculate TFLOPS + flops = 2 * m * n * k # multiply-adds + tflops = flops / (ms * 1e-3) / 1e12 + + # Print shape-specific results with TFLOPS + print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}") + return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/fp8_gemm/", + help="Path to save fp8 gemm benchmark results", + ) + parser.add_argument( + "--run_correctness", + action="store_true", + default=True, + help="Whether to run correctness test", + ) + parser.add_argument( + "--tp_size", + type=int, + default=1, + help="Tensor parallelism size to benchmark (default: 1)", + ) + args = parser.parse_args() + + # Set random seed for reproducibility + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + # Enable TF32, adapted from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_core.py#L148 + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + # Run correctness tests on a few examples + if args.run_correctness: + print("Running correctness tests...") + calculate_diff(64, 512, 7168) # Small test + calculate_diff(64, 7168, 16384) # Medium test + calculate_diff(64, 18432, 7168) # Large test + + # Get the benchmark function with the specified tp_size + benchmark = get_benchmark(args.tp_size) + + print(f"Running performance benchmark for TP size = {args.tp_size}...") + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py new file mode 100644 index 000000000..2c3e8dfcc --- /dev/null +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py @@ -0,0 +1,486 @@ +from typing import Tuple + +import deep_gemm +import torch +import triton +import triton.language as tl +from deep_gemm import calc_diff, get_col_major_tma_aligned_tensor + +# Import shared functionality from the regular GEMM benchmark +from sglang.benchmark.kernels.deepseek.benchmark_deepgemm_fp8_gemm import ( + per_block_cast_to_fp8, + per_token_cast_to_fp8, +) + + +def construct_grouped_and_flat_fp8( + x: torch.Tensor, y: torch.Tensor, num_groups: int, is_masked: bool +) -> Tuple[ + Tuple[torch.Tensor, torch.Tensor], # grouped x_fp8 + Tuple[torch.Tensor, torch.Tensor], # grouped y_fp8 + Tuple[torch.Tensor, torch.Tensor], # flat x_fp8 + Tuple[torch.Tensor, torch.Tensor], # flat y_fp8 + torch.Tensor, # output + torch.Tensor, # reference output +]: + # Verify input shapes + m, k = x.shape + n, k_y = y.shape + assert k == k_y, f"Incompatible shapes: x({m}, {k}), y({n}, {k_y})" + assert m % num_groups == 0, f"m({m}) must be divisible by num_groups({num_groups})" + assert m % 4 == 0, f"TMA alignment error: {m}" + + # Reshape inputs for grouped processing + m_per_group = m // num_groups + x_grouped = x.view(num_groups, m_per_group, k) + y_grouped = y.unsqueeze(0).expand(num_groups, n, k) + + # Initialize output tensors + out = torch.empty((num_groups, m_per_group, n), device="cuda", dtype=torch.bfloat16) + ref_out = torch.einsum("gmk,gnk->gmn", x_grouped, y_grouped) + + # Quantize grouped tensors + x_fp8_grouped = ( + torch.empty_like(x_grouped, dtype=torch.float8_e4m3fn), + torch.empty( + (num_groups, m_per_group, k // 128), device="cuda", dtype=torch.float + ), + ) + y_fp8_grouped = ( + torch.empty_like(y_grouped, dtype=torch.float8_e4m3fn), + torch.empty( + (num_groups, (n + 127) // 128, k // 128), device="cuda", dtype=torch.float + ), + ) + for i in range(num_groups): + x_fp8_grouped[0][i], x_fp8_grouped[1][i] = per_token_cast_to_fp8(x_grouped[i]) + y_fp8_grouped[0][i], y_fp8_grouped[1][i] = per_block_cast_to_fp8(y_grouped[i]) + + # Quantize flat tensors + x_fp8_flat = per_token_cast_to_fp8(x) + y_fp8_flat = per_block_cast_to_fp8(y) + + # For non-masked input, merge the group and M dims in output + if not is_masked: + x_fp8_grouped = ( + x_fp8_grouped[0].view(-1, k), + per_token_cast_to_fp8(x_grouped.view(-1, k))[1], + ) + out, ref_out = out.view(-1, n), ref_out.view(-1, n) + + # Transpose earlier for testing + x_fp8_grouped = ( + x_fp8_grouped[0], + get_col_major_tma_aligned_tensor(x_fp8_grouped[1]), + ) + x_fp8_flat = (x_fp8_flat[0], get_col_major_tma_aligned_tensor(x_fp8_flat[1])) + + return x_fp8_grouped, y_fp8_grouped, x_fp8_flat, y_fp8_flat, out, ref_out + + +# Since we don't have a group gemm kernel in SGLang/vLLM, we implemented a +# custom kernel based on the Triton tutorial. +# https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html +@triton.jit +def fp8_gemm_group_triton_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + # Pointers to scaling factors + a_scale_ptr, + b_scale_ptr, + # Matrix dimensions + M, + N, + K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + # Strides for scaling factors + stride_a_scale_m, + stride_a_scale_k, + stride_b_scale_n, + stride_b_scale_k, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Kernel for computing the matmul C = A x B with FP8 inputs and scaling factors. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + + Note: Block sizes must be multiples of 32 for optimal TMA performance. + """ + # Map program ids to the block of C it should compute + pid_group = tl.program_id(axis=0) # Group ID + pid_n = tl.program_id(axis=1) # N dimension ID + + # Compute the M block ID within this group + group_size_m = min(M - pid_group * GROUP_SIZE_M, GROUP_SIZE_M) + pid_m_within_group = tl.program_id(axis=2) % group_size_m + pid_m = pid_group * GROUP_SIZE_M + pid_m_within_group + + # Create pointers for the first blocks of A and B + 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_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # Initialize accumulator + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + # Main loop + for k_block in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + k_offset = k_block * BLOCK_SIZE_K + + # Load the next block of A and B, with masks + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k_offset, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k_offset, other=0.0) + + # Calculate indices for scaling factors for this K block + a_scale_ptrs = a_scale_ptr + ( + offs_am * stride_a_scale_m + k_block * stride_a_scale_k + ) + b_scale_ptrs = b_scale_ptr + ( + pid_n * stride_b_scale_n + k_block * stride_b_scale_k + ) + + # Perform matrix multiplication in FP8 + res = tl.dot(a, b) + + # Load scaling factors for the current block + a_scale = tl.load(a_scale_ptrs)[:, None] # [BLOCK_SIZE_M, 1] + b_scale = tl.load(b_scale_ptrs) + + # Apply scaling factors to the accumulated result + accumulator += res * a_scale * b_scale + + # Advance pointers + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # Convert to bfloat16 for output + c = accumulator.to(tl.bfloat16) + + # Write back the result + 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_ptr + 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 fp8_gemm_group_triton(a_tuple, b_tuple, c, num_groups): + """ + Perform matrix multiplication with FP8 inputs and proper scaling. + + Args: + a_tuple: Tuple of (quantized_tensor, scale_factors) for input A + b_tuple: Tuple of (quantized_tensor, scale_factors) for input B + c: Output tensor in BF16 format + num_groups: Number of groups for grouped GEMM + + Returns: + Result tensor in BF16 format + """ + # Unpack the tuples + a, a_scale = a_tuple + b, b_scale = b_tuple + + M, K = a.shape + _, N = b.shape + + # Configure block sizes - must be multiples of 32 for TMA alignment + BLOCK_SIZE_M = 128 + BLOCK_SIZE_N = 128 + BLOCK_SIZE_K = 128 + + # Calculate grid dimensions + num_pid_m = triton.cdiv(M, BLOCK_SIZE_M) + num_pid_n = triton.cdiv(N, BLOCK_SIZE_N) + num_groups_grid = triton.cdiv(num_pid_m, num_groups) + + # 3D grid launch - (group, n_blocks, m_blocks_per_group) + grid = (num_groups_grid, num_pid_n, min(num_groups, num_pid_m)) + + fp8_gemm_group_triton_kernel[grid]( + a, + b, + c, + a_scale, + b_scale, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + c.stride(0), + c.stride(1), + a_scale.stride(0), + 1, # Stride in the K dimension may be 1 + b_scale.stride(0), + 1 if b_scale.dim() > 1 else 0, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + GROUP_SIZE_M=num_groups, + ) + + return c + + +def fp8_gemm_group_deepgemm(x_fp8_grouped, y_fp8_grouped, out, m_indices): + deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous( + x_fp8_grouped, + y_fp8_grouped, + out, + m_indices, + ) + return out + + +def calculate_diff(m: int, n: int, k: int, num_groups: int): + print(f"Shape (m={m}, n={n}, k={k}") + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + x_fp8_grouped, y_fp8_grouped, x_fp8_flat, y_fp8_flat, out, out_torch = ( + construct_grouped_and_flat_fp8(x, y, num_groups, is_masked=False) + ) + m_per_group = m // num_groups + out_deepgemm = out.clone() + m_indices = torch.arange(0, num_groups, device="cuda", dtype=torch.int) + m_indices = ( + m_indices.unsqueeze(-1).expand(num_groups, m_per_group).contiguous().view(-1) + ) + + fp8_gemm_group_deepgemm( + x_fp8_grouped, + y_fp8_grouped, + out_deepgemm, + m_indices, + ) + torch.cuda.synchronize() + + # Prepare inputs for Triton + a, a_scale = x_fp8_flat + b, b_scale = y_fp8_flat + b = b.T.contiguous() + # Ensure scales are in the right format and contiguous + a_scale, b_scale = a_scale.contiguous(), b_scale.contiguous() + M, _ = a.shape + _, N = b.shape + c = torch.empty((M, N), device=a.device, dtype=torch.bfloat16) + out_triton = fp8_gemm_group_triton((a, a_scale), (b, b_scale), c, num_groups) + torch.cuda.synchronize() + + diff_torch_deepgemm = torch.abs(out_torch - out_deepgemm).mean().item() + diff_torch_triton = torch.abs(out_torch - out_triton).mean().item() + diff_deepgemm_triton = torch.abs(out_deepgemm - out_triton).mean().item() + + print(f"Shape m={m}, n={n}, k={k}:") + print(f"Torch output: {out_torch[0, 0:5]}") + print(f"DeepGEMM output: {out_deepgemm[0, 0:5]}") + print(f"Triton output: {out_triton[0, 0:5]}") + print(f"Mean absolute difference (Torch-DeepGEMM): {diff_torch_deepgemm}") + print(f"Mean absolute difference (Torch-Triton): {diff_torch_triton}") + print(f"Mean absolute difference (DeepGEMM-Triton): {diff_deepgemm_triton}") + + deepgemm_torch_diff = calc_diff(out_deepgemm, out_torch) + triton_torch_diff = calc_diff(out_triton, out_torch) + deepgemm_triton_diff = calc_diff(out_deepgemm, out_triton) + + DIFF_THRESHOLD = 0.001 + all_match = ( + deepgemm_torch_diff < DIFF_THRESHOLD + and triton_torch_diff < DIFF_THRESHOLD + and deepgemm_triton_diff < DIFF_THRESHOLD + ) + if all_match: + print("✅ All implementations match\n") + else: + print("❌ Some implementations differ:") + print( + f" - Torch vs DeepGEMM: {'✅' if deepgemm_torch_diff < DIFF_THRESHOLD else '❌'}" + f" - Torch vs Triton: {'✅' if triton_torch_diff < DIFF_THRESHOLD else '❌'}" + f" - DeepGEMM vs Triton: {'✅' if deepgemm_triton_diff < DIFF_THRESHOLD else '❌'}" + ) + + +def get_weight_shapes(tp_size): + # cannot TP + total = [ + (512 + 64, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (7168, 16384), + (7168, 18432), + ] + # N can TP + n_tp = [ + (18432 * 2, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (24576, 1536), + (4096, 7168), + ] + # K can TP + k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] + + weight_shapes = [] + for t in total: + weight_shapes.append(t) + for n_t in n_tp: + new_t = (n_t[0] // tp_size, n_t[1]) + weight_shapes.append(new_t) + for k_t in k_tp: + new_t = (k_t[0], k_t[1] // tp_size) + weight_shapes.append(new_t) + + return weight_shapes + + +def create_benchmark_configs(tp_size): + configs = [] + weight_shapes = get_weight_shapes(tp_size) + batch_sizes = [2048, 4096] + group_sizes = [4, 8] + for n, k in weight_shapes: + for m in batch_sizes: + for num_groups in group_sizes: + configs.append((m, n, k, num_groups, tp_size)) + + return configs + + +def get_benchmark(tp_size): + all_configs = create_benchmark_configs(tp_size) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["m", "n", "k", "num_groups", "tp_size"], + x_vals=[config for config in all_configs], + line_arg="provider", + line_vals=["deepgemm", "triton"], + line_names=["DeepGEMM", "Triton"], + styles=[("blue", "-"), ("red", "-")], + ylabel="ms", + plot_name=f"fp8-group-gemm-performance-comparison-tp{tp_size}", + args={}, + ) + ) + def benchmark(m, n, k, num_groups, tp_size, provider): + print( + f"Shape (m={m}, n={n}, k={k}, tp={tp_size}, num_groups={num_groups}, Provider: {provider}" + ) + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + x_fp8_grouped, y_fp8_grouped, x_fp8_flat, y_fp8_flat, out, out_torch = ( + construct_grouped_and_flat_fp8(x, y, num_groups, is_masked=False) + ) + m_per_group = m // num_groups + m_indices = torch.arange(0, num_groups, device="cuda", dtype=torch.int) + m_indices = ( + m_indices.unsqueeze(-1) + .expand(num_groups, m_per_group) + .contiguous() + .view(-1) + ) + + quantiles = [0.5, 0.2, 0.8] + + if provider == "deepgemm": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: fp8_gemm_group_deepgemm( + x_fp8_grouped, + y_fp8_grouped, + out, + m_indices, + ), + quantiles=quantiles, + ) + elif provider == "triton": + # Prepare inputs for Triton + # We did it outside of the lambda function to make it fair comparison like deepgemm + a, a_scale = x_fp8_flat + b, b_scale = y_fp8_flat + b = b.T.contiguous() + # Ensure scales are in the right format and contiguous + a_scale, b_scale = a_scale.contiguous(), b_scale.contiguous() + M, _ = a.shape + _, N = b.shape + c = torch.empty((M, N), device=a.device, dtype=torch.bfloat16) + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: fp8_gemm_group_triton( + (a, a_scale), + (b, b_scale), + c, + num_groups, + ), + quantiles=quantiles, + ) + + # Calculate TFLOPS + flops = 2 * m * n * k # multiply-adds + tflops = flops / (ms * 1e-3) / 1e12 + + print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}") + return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/fp8_group_gemm/", + help="Path to save deepgemm fp8 group gemm benchmark results", + ) + parser.add_argument( + "--run_correctness", + action="store_true", + help="Whether to run correctness test", + ) + parser.add_argument( + "--tp_size", + type=int, + default=1, + help="Tensor parallelism size to benchmark (default: 1)", + ) + args = parser.parse_args() + + # Set random seed for reproducibility + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + # Enable TF32, adapted from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_core.py#L148 + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + # Run correctness tests on a few examples + if args.run_correctness: + print("Running correctness tests...") + calculate_diff(8192, 7168, 4096, 4) + calculate_diff(8192, 2048, 7168, 4) + calculate_diff(4096, 7168, 4096, 8) + calculate_diff(4096, 2048, 7168, 8) + calculate_diff(4096, 576, 7168, 8) + + # Get the benchmark function with the specified tp_size + benchmark = get_benchmark(args.tp_size) + + print(f"Running performance benchmark for TP size = {args.tp_size}...") + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/benchmark/kernels/fbgemm/README.md b/benchmark/kernels/fbgemm/README.md new file mode 100644 index 000000000..e51356d8a --- /dev/null +++ b/benchmark/kernels/fbgemm/README.md @@ -0,0 +1,29 @@ +## Benchmark FBGEMM Grouped GEMM + +Benchmark FBGEMM Grouped GEMM in both Triton and CUDA version and SGLang Triton Grouped GEMM, it will be used to compare the bandwidth of different implementations. + +### Requirements + +```shell +pip install fbgemm-gpu-genai +``` + +### Usage + +```bash +python3 benchmark/fbgemm/benchmark_fbgemm_grouped_gemm.py --model Qwen/Qwen2-57B-A14B-Instruct --tp-size 4 --use-fp8-w8a8 +``` + +For example, in H200, the Qwen2-57B-A14B-Instruct TP4 fp8w8a8 grouped gemm bandwidth result is as follows: + +```shell +grouped-gemm-performance: + batch_size FBGEMM Triton Grouped GEMM FP8 FBGEMM CUTLASS F8F8BF16 Rowwise SGLang Grouped GEMM FP8 +0 256.0 3704.841339 3042.626402 2254.725030 +1 512.0 3691.426346 3029.065684 2269.504543 +2 1024.0 3653.938629 2258.471467 2358.319020 +3 2048.0 3596.644313 2271.611904 2476.895397 +4 4096.0 3468.496435 2231.283986 2179.473910 +``` + +The theoretical peak bandwidth of H200 is 4.8 TB/s. Taking batch_size 256 as an example, the bandwidth of FBGEMM Triton Grouped GEMM FP8 is 3704.841339 GB/s, the bandwidth of FBGEMM CUTLASS F8F8BF16 Rowwise is 3042.626402 GB/s, and the bandwidth of SGLang Grouped GEMM FP8 is 2254.725030 GB/s. Therefore, FBGEMM Triton Grouped GEMM FP8 achieves 77.9% of H200's theoretical peak bandwidth, FBGEMM CUTLASS F8F8BF16 Rowwise achieves 63.4% of H200's theoretical peak bandwidth, and SGLang Grouped GEMM FP8 achieves 46.9% of H200's theoretical peak bandwidth. diff --git a/benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py b/benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py new file mode 100644 index 000000000..6e8c8dcf2 --- /dev/null +++ b/benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py @@ -0,0 +1,516 @@ +# python3 benchmark/fbgemm/benchmark_fbgemm_grouped_gemm.py --model Qwen/Qwen2-57B-A14B-Instruct --tp-size 4 --use-fp8-w8a8 +import argparse + +import torch +import triton +from fbgemm_gpu.experimental.gemm.triton_gemm.fp8_gemm import ( + quantize_fp8_row, + triton_quantize_fp8_row, +) +from fbgemm_gpu.experimental.gemm.triton_gemm.grouped_gemm import ( + grouped_gemm as fbgemm_grouped_gemm, +) +from fbgemm_gpu.experimental.gemm.triton_gemm.grouped_gemm import ( + grouped_gemm_fp8_rowwise as fbgemm_grouped_gemm_fp8_rowwise, +) +from transformers import AutoConfig + +from sglang.srt.layers.moe.ep_moe.kernels import ( + grouped_gemm_triton as sglang_grouped_gemm, +) + + +def get_model_config(model_name: str, tp_size: int): + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + + if config.architectures[0] == "DbrxForCausalLM": + num_groups = config.ffn_config.moe_num_experts + intermediate_size = config.ffn_config.ffn_hidden_size + elif config.architectures[0] == "JambaForCausalLM": + num_groups = config.num_experts + intermediate_size = config.intermediate_size + elif config.architectures[0] == "Qwen2MoeForCausalLM": + num_groups = config.num_experts + intermediate_size = config.moe_intermediate_size + elif config.architectures[0] == "Qwen3MoeForCausalLM": + num_groups = config.num_experts + intermediate_size = config.moe_intermediate_size + elif config.architectures[0] in [ + "DeepseekV2ForCausalLM", + "DeepseekV3ForCausalLM", + ]: + num_groups = config.n_routed_experts + intermediate_size = config.moe_intermediate_size + elif config.architectures[0] == "Llama4ForConditionalGeneration": + num_groups = config.text_config.num_local_experts + intermediate_size = config.text_config.intermediate_size + elif config.architectures[0] in [ + "Grok1ForCausalLM", + "Grok1ImgGen", + "Grok1AForCausalLM", + ]: + num_groups = config.num_local_experts + intermediate_size = config.moe_intermediate_size + else: + num_groups = config.num_local_experts + intermediate_size = config.intermediate_size + + shape_configs = { + "num_groups": num_groups, + "hidden_size": config.hidden_size, + "intermediate_size": intermediate_size, + "dtype": config.torch_dtype, + } + print(f"{shape_configs=}") + return shape_configs + + +def create_test_data(batch_size, num_groups, hidden_size, intermediate_size): + torch.manual_seed(42) + + tokens_per_group = batch_size // num_groups + m_sizes = torch.full( + (num_groups,), tokens_per_group, dtype=torch.int32, device="cuda" + ) + + x = torch.randn(batch_size, hidden_size, dtype=torch.bfloat16, device="cuda") + + base_weights = torch.randn( + num_groups, intermediate_size, hidden_size, dtype=torch.bfloat16, device="cuda" + ) + + w_fbgemm = base_weights.reshape(num_groups * intermediate_size, hidden_size) + w_sglang = base_weights + + c_fbgemm = torch.empty( + batch_size, intermediate_size, dtype=torch.bfloat16, device="cuda" + ) + c_sglang = torch.empty( + batch_size, intermediate_size, dtype=torch.bfloat16, device="cuda" + ) + + seg_indptr = torch.zeros(num_groups + 1, dtype=torch.int32, device="cuda") + for i in range(1, num_groups + 1): + seg_indptr[i] = seg_indptr[i - 1] + tokens_per_group + + weight_indices = torch.arange(num_groups, dtype=torch.int32, device="cuda") + + return ( + x, + w_fbgemm, + w_sglang, + c_fbgemm, + c_sglang, + m_sizes, + seg_indptr, + weight_indices, + ) + + +def create_fp8_test_data( + batch_size, num_groups, hidden_size, intermediate_size, backend="triton" +): + """ + Create test data for FP8 grouped GEMM operations. + + Args: + batch_size: Total batch size + num_groups: Number of groups + hidden_size: Hidden dimension size + intermediate_size: Intermediate dimension size + backend: "triton" for Triton GEMM, "cutlass" for CUTLASS GEMM + + Returns: + For triton: (x_fp8, w_fp8, m_sizes, x_scale, w_scale) + For cutlass: (x, wq, w_scale, m_sizes) + """ + torch.manual_seed(42) + + tokens_per_group = batch_size // num_groups + + # Create weight matrices for each group + w_list = [] + for _ in range(num_groups): + w = torch.randn( + intermediate_size, hidden_size, dtype=torch.float16, device="cuda" + ) + w_list.append(w) + + # Quantize weights using quantize_fp8_row for each group + wq_list, w_scale_list = zip(*[quantize_fp8_row(w) for w in w_list]) + + if backend == "triton": + # Triton format: concatenated weights + w_fp8 = torch.concat(wq_list, dim=0).contiguous() + w_scale = torch.concat(w_scale_list, dim=0).contiguous() + + # Create m_sizes as int32 for triton + m_sizes = torch.full( + (num_groups,), tokens_per_group, dtype=torch.int32, device="cuda" + ) + + # Create and quantize input + x_fp16 = torch.randn( + batch_size, hidden_size, dtype=torch.float16, device="cuda" + ) + x_fp8, x_scale = triton_quantize_fp8_row(x_fp16) + x_scale = x_scale.view(batch_size, -1) + + return x_fp8, w_fp8, m_sizes, x_scale, w_scale + + elif backend == "cutlass": + # CUTLASS format: stacked weights + wq = torch.stack(wq_list, dim=0).contiguous() + w_scale = torch.stack(w_scale_list, dim=0).contiguous() + + # Create m_sizes as int64 for cutlass + m_values = [tokens_per_group] * num_groups + m_sizes = torch.tensor(m_values).to(dtype=torch.int64, device="cuda") + + # Create input data - separate for each group then concat + x_list = [] + for _ in range(num_groups): + x = torch.randn( + tokens_per_group, hidden_size, dtype=torch.float16, device="cuda" + ) + x_list.append(x) + + # Concatenate inputs into single tensor + x = torch.concat(x_list, dim=0).contiguous() + + return x, wq, w_scale, m_sizes + + else: + raise ValueError(f"Unsupported backend: {backend}") + + +def calculate_memory_bandwidth(m_sizes, hidden_size, intermediate_size, dtype): + """ + Calculate memory bandwidth based on accessed expert weights. + + Args: + m_sizes: Tensor containing batch sizes for each group + hidden_size: Hidden dimension size + intermediate_size: Intermediate dimension size + dtype: Data type of weights + + Returns: + Memory size in bytes for accessed expert weights + """ + # Count non-zero groups (active experts) + if hasattr(m_sizes, "cpu"): + active_experts = torch.count_nonzero(m_sizes).item() + else: + active_experts = sum(1 for m in m_sizes if m > 0) + + # Calculate bytes per element based on dtype + if dtype in [torch.float16, torch.bfloat16]: + bytes_per_element = 2 + elif dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + bytes_per_element = 1 + elif dtype == torch.float32: + bytes_per_element = 4 + else: + # Default to 2 bytes for unknown dtypes + bytes_per_element = 2 + + # Memory per expert weight matrix + memory_per_expert = hidden_size * intermediate_size * bytes_per_element + + # Total memory for active experts + total_memory_bytes = active_experts * memory_per_expert + + return total_memory_bytes + + +def get_benchmark_config(use_fp8_w8a8=False): + if use_fp8_w8a8: + return { + "line_vals": [ + "fbgemm_triton_grouped_gemm_fp8", + "fbgemm_cutlass_f8f8bf16_rowwise", + "sglang_grouped_gemm", + ], + "line_names": [ + "FBGEMM Triton Grouped GEMM FP8", + "FBGEMM CUTLASS F8F8BF16 Rowwise", + "SGLang Grouped GEMM FP8", + ], + "styles": [("blue", "-"), ("orange", "-"), ("red", "-")], + } + else: + return { + "line_vals": ["fbgemm_triton_grouped_gemm", "sglang_grouped_gemm"], + "line_names": [ + "FBGEMM Triton Grouped GEMM BF16", + "SGLang Grouped GEMM BF16", + ], + "styles": [("blue", "-"), ("green", "-")], + } + + +def run_benchmark( + model_config, use_fp8_w8a8=False, save_path="./benchmark_grouped_gemm/" +): + config = get_benchmark_config(use_fp8_w8a8) + + benchmark_config = triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=[256, 512, 1024, 2048, 4096], + line_arg="provider", + line_vals=config["line_vals"], + line_names=config["line_names"], + styles=config["styles"], + ylabel="Bandwidth (GB/s)", + plot_name="grouped-gemm-performance", + args={}, + ) + + @triton.testing.perf_report(benchmark_config) + def dynamic_benchmark(batch_size, provider, model_config, use_fp8_w8a8=False): + print(f"Benchmarking {provider} with batch_size={batch_size}") + torch.cuda.manual_seed_all(0) + + num_groups = model_config["num_groups"] + hidden_size = model_config["hidden_size"] + intermediate_size = model_config["intermediate_size"] + + if provider == "fbgemm_triton_grouped_gemm_fp8": + try: + test_data = create_fp8_test_data( + batch_size, + num_groups, + hidden_size, + intermediate_size, + backend="triton", + ) + x_fp8, w_fp8, m_sizes, x_scale, w_scale = test_data + + # Calculate memory bandwidth + memory_bytes = calculate_memory_bandwidth( + m_sizes, hidden_size, intermediate_size, torch.float8_e4m3fn + ) + + def run_func(): + return fbgemm_grouped_gemm_fp8_rowwise( + x_fp8, w_fp8, m_sizes, x_scale, w_scale, use_fast_accum=True + ) + + except Exception as e: + print(f"FP8 not supported, skipping: {e}") + return float("inf"), float("inf"), float("inf") + + elif provider == "fbgemm_cutlass_f8f8bf16_rowwise": + try: + test_data = create_fp8_test_data( + batch_size, + num_groups, + hidden_size, + intermediate_size, + backend="cutlass", + ) + x, wq, w_scale, m_sizes = test_data + + # Calculate memory bandwidth + memory_bytes = calculate_memory_bandwidth( + m_sizes, hidden_size, intermediate_size, torch.float8_e4m3fn + ) + + # Quantize input using triton_quantize_fp8_row + xq, x_scale = triton_quantize_fp8_row(x) + x_scale = x_scale.view(batch_size, -1) + + def run_func(): + return torch.ops.fbgemm.f8f8bf16_rowwise_grouped_stacked( + xq, wq, x_scale, w_scale, m_sizes + ) + + except Exception as e: + print( + f"CUTLASS f8f8bf16_rowwise_grouped_stacked not supported, " + f"skipping: {e}" + ) + return float("inf"), float("inf"), float("inf") + else: + test_data = create_test_data( + batch_size, num_groups, hidden_size, intermediate_size + ) + ( + x, + w_fbgemm, + w_sglang, + c_fbgemm, + c_sglang, + m_sizes, + seg_indptr, + weight_indices, + ) = test_data + + # Calculate memory bandwidth for BF16 operations + memory_bytes = calculate_memory_bandwidth( + m_sizes, hidden_size, intermediate_size, torch.bfloat16 + ) + + if provider == "fbgemm_triton_grouped_gemm": + + def run_func(): + return fbgemm_grouped_gemm( + x, w_fbgemm, m_sizes, use_fast_accum=True + ) + + else: + + def run_func(): + return sglang_grouped_gemm( + x, + w_sglang, + c_sglang, + num_groups, + weight_column_major=True, + seg_indptr=seg_indptr, + weight_indices=weight_indices, + c_dtype=c_sglang.dtype, + ) + + for _ in range(10): + try: + run_func() + except Exception as e: + print(f"Error during warmup for {provider}: {e}") + return float("inf"), float("inf"), float("inf") + + torch.cuda.synchronize() + + try: + quantiles = [0.5, 0.2, 0.8] + ms, min_ms, max_ms = triton.testing.do_bench(run_func, quantiles=quantiles) + + # Convert time (ms) to bandwidth (GB/s) + # Bandwidth = Memory (bytes) / Time (seconds) + # Convert ms to seconds and bytes to GB (1e9) + gb_per_s = (memory_bytes / 1e9) / (ms / 1000) + # min bandwidth = max time, max bandwidth = min time + min_gb_per_s = (memory_bytes / 1e9) / (max_ms / 1000) + max_gb_per_s = (memory_bytes / 1e9) / (min_ms / 1000) + + return gb_per_s, min_gb_per_s, max_gb_per_s + except Exception as e: + print(f"Error during benchmarking for {provider}: {e}") + return 0.0, 0.0, 0.0 + + dynamic_benchmark.run( + show_plots=True, + print_data=True, + save_path=save_path, + model_config=model_config, + use_fp8_w8a8=use_fp8_w8a8, + ) + + +def verify_correctness(model_config): + print("Verifying correctness...") + batch_size = 128 + num_groups = model_config["num_groups"] + hidden_size = model_config["hidden_size"] + intermediate_size = model_config["intermediate_size"] + + test_data = create_test_data(batch_size, num_groups, hidden_size, intermediate_size) + ( + x, + w_fbgemm, + w_sglang, + c_fbgemm, + c_sglang, + m_sizes, + seg_indptr, + weight_indices, + ) = test_data + + result_fbgemm = fbgemm_grouped_gemm(x, w_fbgemm, m_sizes, use_fast_accum=True) + + result_sglang = sglang_grouped_gemm( + x, + w_sglang, + c_sglang, + num_groups, + weight_column_major=True, + seg_indptr=seg_indptr, + weight_indices=weight_indices, + c_dtype=c_sglang.dtype, + ) + + if torch.allclose(result_fbgemm, result_sglang, rtol=1e-3, atol=1e-3): + print("✓ BF16 Correctness verification passed!") + else: + max_diff = torch.max(torch.abs(result_fbgemm - result_sglang)) + print(f"✗ BF16 Correctness verification failed! Max diff: {max_diff}") + return False + + return True + + +def main(): + parser = argparse.ArgumentParser( + description="Benchmark FBGEMM vs SGLang Grouped GEMM" + ) + parser.add_argument( + "--model", + type=str, + default="mistralai/Mixtral-8x7B-Instruct-v0.1", + help="Model name to get configuration from", + ) + parser.add_argument( + "--tp-size", type=int, default=1, help="Tensor parallelism size" + ) + parser.add_argument( + "--use-fp8-w8a8", action="store_true", help="Enable FP8 W8A8 benchmark" + ) + parser.add_argument( + "--save-path", + type=str, + default="./benchmark_grouped_gemm/", + help="Path to save benchmark results", + ) + parser.add_argument( + "--verify-correctness", + action="store_true", + help="Verify correctness before benchmarking", + ) + + args = parser.parse_args() + + try: + model_config = get_model_config(args.model, args.tp_size) + except Exception as e: + print(f"Failed to get model config: {e}") + print("Using default configuration...") + model_config = { + "num_groups": 8, + "hidden_size": 4096, + "intermediate_size": 14336, + "dtype": torch.bfloat16, + } + + print("Running benchmark with:") + print(f" num_groups: {model_config['num_groups']}") + print(f" hidden_size: {model_config['hidden_size']}") + print(f" intermediate_size: {model_config['intermediate_size']}") + print(f" use_fp8_w8a8: {args.use_fp8_w8a8}") + + if args.verify_correctness: + if not verify_correctness(model_config): + print("Correctness verification failed. Exiting...") + return + + try: + run_benchmark( + model_config=model_config, + use_fp8_w8a8=args.use_fp8_w8a8, + save_path=args.save_path, + ) + except Exception as e: + print(f"Benchmark failed: {e}") + + +if __name__ == "__main__": + main() diff --git a/benchmark/kernels/flashinfer_allreduce_fusion/README.md b/benchmark/kernels/flashinfer_allreduce_fusion/README.md new file mode 100644 index 000000000..e651604c7 --- /dev/null +++ b/benchmark/kernels/flashinfer_allreduce_fusion/README.md @@ -0,0 +1,102 @@ +# FlashInfer Fused AllReduce + RMSNorm Benchmark + +This benchmark script is modified from the [original implementation](https://github.com/vllm-project/vllm/blob/237e1fb887c7f5a579420fa0295097f24b006594/benchmarks/kernels/benchmark_fused_collective.py) by the vLLM community. It aims to compare the performance differences between FlashInfer fused operators in SGLang (trtllm_allreduce_fusion: AllReduce + Residual Add + RMSNorm + optional quantization) and conventional implementations (standard `tensor_model_parallel_all_reduce` + separate RMSNorm/quantization). Specifically, this script tests the timing performance of two implementation paths: 1) Standard AllReduce and RMSNorm executed separately; 2) FlashInfer's fused operator combining AllReduce, Residual Add, RMSNorm, and optional quantization operations. + +This benchmark script helps us tune the ipc workspace size of the `flashinfer_allreduce_residual_rmsnorm` operator in SGLang and prepare for applications with FP8/FP4 quantized fused operators. + +Script path: `benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py` + +## Feature Overview + +- Compare average execution time (ms) and calculate speedup ratios for the following paths: + - standard_allreduce_rmsnorm (Standard AllReduce + RMSNorm) + - flashinfer_fused_allreduce_rmsnorm (Fused AllReduce + RMSNorm), including oneshot and twoshot modes + - Optionally compare FP8/FP4 quantized fused paths with standard paths +- Use CUDA Graph capture and batch replay to reduce measurement noise +- Automatically select the faster "standard baseline" (native/compiled version) as the denominator for speedup calculation +- Optionally export results in Markdown format + +## Runtime Environment and Prerequisites + +- At least 2 GPUs, and launch multi-process distributed training using `torchrun` (NCCL backend) +- Properly install/compile sglang along with sgl-kernel and custom operators + +## Quick Start (Command Examples) + +The following examples use world_size=2. You can modify `--nproc_per_node` and parameters according to your machine: + +- Regular paths only (no quantization): +``` +torchrun --nproc_per_node=2 \ +benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py \ +--no-quant --hidden-dim 1024 --seq-lens 512 1024 2048 4096 --trials 100 +``` + +- FP8 quantization paths only: +``` +torchrun --nproc_per_node=2 \ +benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py \ +--quant-fp8 --hidden-dim 1024 --seq-lens 512 1024 2048 4096 --trials 100 +``` + +- FP4 quantization paths only: +``` +torchrun --nproc_per_node=2 \ +benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py \ +--quant-fp4 --hidden-dim 1024 --seq-lens 512 1024 2048 4096 --trials 100 +``` + +- Larger hidden dimensions: +``` +torchrun --nproc_per_node=2 \ +benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py \ +--no-quant --hidden-dim 4096 --seq-lens 512 1024 2048 4096 --trials 100 +``` + +## Parameter Description +- `--seq-lens`: List of sequence lengths to test (default: 128 512 1024 2048) +- `--hidden-dim`: Hidden dimension (default: 8192) +- `--dtypes`: Data type list, `float16|bfloat16|float32` (default: bfloat16) +- `--no-residual`: Only test "no residual" scenarios (default tests both "with/without residual") +- Mutually exclusive quantization options: + - `--no-quant`: No quantization testing + - `--quant-fp8`: Only FP8 quantization testing + - `--quant-fp4`: Only FP4 quantization testing + - `--quant-all`: Test all (default) +- FlashInfer related: + - `--disable-oneshot`: Disable oneshot mode (default enables oneshot and tests twoshot simultaneously) +- Runtime configuration: + - `--warmup`: Warmup count before graph capture and before graph replay (default 5) + - `--trials`: Benchmark iteration count (default 20; internally each `graph.replay()` will batch replay multiple times) + - `--output-file`: Save results as Markdown file (only rank0 takes effect) + +## Output Example + +Each configuration group prints a table showing average execution time and relative speedup ratios (baseline is the faster standard implementation). For example: +``` +================================================================================ +Results: seq_len=1024, hidden_dim=1024 +dtype=torch.bfloat16, residual=yes, quant_mode=none +================================================================================ +Operation Time (ms) Speedup +-------------------------------------------------------------------------------- +standard_allreduce_rmsnorm 0.024 0.98x +standard_allreduce_rmsnorm_native_compiled 0.023 baseline +flashinfer_fused_allreduce_rmsnorm_oneshot 0.011 2.19x +flashinfer_fused_allreduce_rmsnorm_twoshot 0.041 0.57x +``` + +If `--output-file` is specified, all configurations will be summarized in Markdown tables in that file. + +## Important Notes and Recommendations + +- Distributed: The script uses `torchrun` environment variables to initialize distributed training and binds tensors/communication groups to the current rank's corresponding device. +- World size: Requires `WORLD_SIZE > 1` to perform communication operator benchmarks. Otherwise, the script will error and prompt. +- FlashInfer: + - If not installed or interfaces are missing, the script will only run standard paths and provide prompts in the logs. + - The fused operator internally uses "oneshot"/"twoshot" two trigger methods; oneshot is enabled by default and twoshot is tested simultaneously. +- FP8/FP4: + - FP8 uses sglang's FP8 tools and dtype, with underlying platform selection of `e4m3`/`e4m3fnuz` etc. + - FP4 uses sgl-kernel's `scaled_fp4_quant`, requiring corresponding platform support. +- CUDA Graph: + - Uses sglang's `graph_capture()` to prepare capture-ready state for communication, then uses `torch.cuda.graph` to capture kernels, reducing measurement jitter. diff --git a/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py b/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py new file mode 100644 index 000000000..4aebf62b9 --- /dev/null +++ b/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py @@ -0,0 +1,1304 @@ +# Modified from https://github.com/vllm-project/vllm/blob/237e1fb887c7f5a579420fa0295097f24b006594/benchmarks/kernels/benchmark_fused_collective.py + +""" +Benchmark for FlashInfer fused collective operations vs standard operations. + +This benchmark compares: +1. FlashInfer's trtllm_allreduce_fusion (fused allreduce + rmsnorm + optional quant) +2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations + +Usage with torchrun: + torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --no-quant --hidden-dim 1024 --seq-len 512 1024 2048 4096 --trials 100 + torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --quant-fp8 --hidden-dim 1024 --seq-len 512 1024 2048 4096 --trials 100 + torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --quant-fp4 --hidden-dim 1024 --seq-len 512 1024 2048 4096 --trials 100 + + torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --no-quant --hidden-dim 4096 --seq-len 512 1024 2048 4096 --trials 100 + torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --quant-fp8 --hidden-dim 4096 --seq-len 512 1024 2048 4096 --trials 100 + torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --quant-fp4 --hidden-dim 4096 --seq-len 512 1024 2048 4096 --trials 100 +""" + +import argparse +import contextlib +import itertools +import logging +import os +import time +from typing import Optional + +import torch # type: ignore +import torch.distributed as dist # type: ignore + +from sglang.srt.distributed import get_tp_group, tensor_model_parallel_all_reduce +from sglang.srt.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + graph_capture, + init_distributed_environment, + initialize_model_parallel, +) +from sglang.srt.layers.layernorm import RMSNorm # noqa +from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype as SGLANG_FP8_DTYPE +from sglang.srt.layers.quantization.fp8_kernel import static_quant_fp8 + +try: + from sgl_kernel import fused_add_rmsnorm as SGL_FUSED_ADD_RMS_NORM + from sgl_kernel import rmsnorm as SGL_RMS_NORM + from sgl_kernel import scaled_fp4_quant as SGL_SCALED_FP4_QUANT +except Exception: # pragma: no cover - fallback on non-supported platforms + SGL_FUSED_ADD_RMS_NORM = None + SGL_RMS_NORM = None + SGL_SCALED_FP4_QUANT = None + +FP8_DTYPE = SGLANG_FP8_DTYPE + +logger = logging.getLogger(__name__) + +# Try to import FlashInfer +try: + import flashinfer.comm as flashinfer_comm # type: ignore + + if not hasattr(flashinfer_comm, "trtllm_allreduce_fusion"): + flashinfer_comm = None + logger.warning( + "FlashInfer comm module found but missing trtllm_allreduce_fusion" + ) +except ImportError: + flashinfer_comm = None + logger.warning("FlashInfer not found, only benchmarking standard operations") + +# Constants +MiB = 1024 * 1024 + +# FlashInfer max sizes per world size +# Enable 64MB for 2, 4, 8 world sizes to verify large input sizes +# use --disable-oneshot to disable oneshot mode for very large input sizes +_FI_MAX_SIZES = { + 2: 64 * MiB, # 64MB + 4: 64 * MiB, # 64MB + 8: 64 * MiB, # 64MB +} + +# Global workspace tensor for FlashInfer +_FI_WORKSPACE_TENSOR = None + + +def setup_flashinfer_workspace( + world_size: int, + rank: int, + hidden_dim: int, + max_token_num: int, + use_fp32_lamport: bool = False, +): + """Setup FlashInfer workspace for fused allreduce operations.""" + global _FI_WORKSPACE_TENSOR + + if flashinfer_comm is None: + return None, None + + if world_size not in _FI_MAX_SIZES: + logger.warning("FlashInfer not supported for world size %s", world_size) + return None, None + + try: + # Create IPC workspace + ipc_handles, workspace_tensor = ( + flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion( + tp_rank=rank, + tp_size=world_size, + max_token_num=max_token_num, + hidden_dim=hidden_dim, + group=get_tp_group().device_group, + use_fp32_lamport=use_fp32_lamport, + ) + ) + + _FI_WORKSPACE_TENSOR = workspace_tensor + return ipc_handles, workspace_tensor + except Exception as e: + logger.error("Failed to setup FlashInfer workspace: %s", e) + return None, None + + +def cleanup_flashinfer_workspace(ipc_handles): + """Cleanup FlashInfer workspace.""" + if flashinfer_comm is None or ipc_handles is None: + return + + try: + group = get_tp_group().device_group + flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(ipc_handles, group) + except Exception as e: + logger.error("Failed to cleanup FlashInfer workspace: %s", e) + + +class FlashInferFusedAllReduceParams: + """Parameters for FlashInfer fused allreduce operations.""" + + def __init__( + self, + rank: int, + world_size: int, + use_fp32_lamport: bool = False, + max_token_num: int = 1024, + ): + self.rank = rank + self.world_size = world_size + self.use_fp32_lamport = use_fp32_lamport + self.trigger_completion_at_end = True + self.launch_with_pdl = True + self.fp32_acc = True + self.max_token_num = max_token_num + + def get_trtllm_fused_allreduce_kwargs(self): + return { + "world_rank": self.rank, + "world_size": self.world_size, + "launch_with_pdl": self.launch_with_pdl, + "trigger_completion_at_end": self.trigger_completion_at_end, + "fp32_acc": self.fp32_acc, + } + + +def flashinfer_fused_allreduce_rmsnorm( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rms_gamma: torch.Tensor, + rms_eps: float, + allreduce_params: "FlashInferFusedAllReduceParams", + use_oneshot: bool, + norm_out: Optional[torch.Tensor] = None, +): + """FlashInfer fused allreduce + rmsnorm operation.""" + if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None: + raise RuntimeError("FlashInfer not available or workspace not initialized") + + if norm_out is None: + norm_out = input_tensor + residual_out = residual + else: + residual_out = input_tensor + + flashinfer_comm.trtllm_allreduce_fusion( + allreduce_in=input_tensor, + token_num=input_tensor.shape[0], + residual_in=residual, + residual_out=residual_out, + norm_out=norm_out, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + hidden_dim=input_tensor.shape[-1], + workspace_ptrs=_FI_WORKSPACE_TENSOR, + pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, + allreduce_out=None, + quant_out=None, + scale_out=None, + layout_code=None, + scale_factor=None, + use_oneshot=use_oneshot, + **allreduce_params.get_trtllm_fused_allreduce_kwargs(), + ) + + +def flashinfer_fused_allreduce_rmsnorm_fp8_quant( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rms_gamma: torch.Tensor, + rms_eps: float, + scale_factor: torch.Tensor, + allreduce_params: FlashInferFusedAllReduceParams, + use_oneshot: bool = True, + norm_out: Optional[torch.Tensor] = None, + quant_out: Optional[torch.Tensor] = None, +): + """FlashInfer fused allreduce + rmsnorm + FP8 quantization.""" + if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None: + raise RuntimeError("FlashInfer not available or workspace not initialized") + + if norm_out is None: + norm_out = input_tensor + residual_out = residual + else: + residual_out = input_tensor + + flashinfer_comm.trtllm_allreduce_fusion( + allreduce_in=input_tensor, + token_num=input_tensor.shape[0], + residual_in=residual, + residual_out=residual_out, + norm_out=norm_out, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + hidden_dim=input_tensor.shape[-1], + workspace_ptrs=_FI_WORKSPACE_TENSOR, + pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + allreduce_out=None, + quant_out=quant_out, + scale_out=None, + layout_code=None, + scale_factor=scale_factor, + use_oneshot=use_oneshot, + **allreduce_params.get_trtllm_fused_allreduce_kwargs(), + ) + + +def flashinfer_fused_allreduce_rmsnorm_fp4_quant( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rms_gamma: torch.Tensor, + rms_eps: float, + input_global_scale: torch.Tensor, + allreduce_params: FlashInferFusedAllReduceParams, + quant_out: torch.Tensor, + use_oneshot: bool, + output_scale: torch.Tensor, + norm_out: Optional[torch.Tensor] = None, +): + """FlashInfer fused allreduce + rmsnorm + FP4 quantization.""" + if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None: + raise RuntimeError("FlashInfer not available or workspace not initialized") + + if norm_out is None: + norm_out = input_tensor + residual_out = residual + else: + residual_out = input_tensor + + flashinfer_comm.trtllm_allreduce_fusion( + allreduce_in=input_tensor, + token_num=input_tensor.shape[0], + residual_in=residual, + residual_out=residual_out, + norm_out=norm_out, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + hidden_dim=input_tensor.shape[-1], + workspace_ptrs=_FI_WORKSPACE_TENSOR, + pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + allreduce_out=None, + quant_out=quant_out, + scale_out=output_scale, + layout_code=None, + scale_factor=input_global_scale, + use_oneshot=use_oneshot, + **allreduce_params.get_trtllm_fused_allreduce_kwargs(), + ) + + +def standard_allreduce_rmsnorm( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rms_gamma: torch.Tensor, + rms_eps: float, + norm_out: Optional[torch.Tensor] = None, +): + """Standard allreduce + rmsnorm operations.""" + # All-reduce first + allreduce_out = tensor_model_parallel_all_reduce(input_tensor) + # Then RMS norm + if residual is not None: + # Fused add + RMS norm (in-place on allreduce_out) + if SGL_FUSED_ADD_RMS_NORM is not None: + SGL_FUSED_ADD_RMS_NORM(allreduce_out, residual, rms_gamma, rms_eps) + else: + rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) + rms.weight.data = rms_gamma + rms.forward_native(allreduce_out, residual) + else: + # Just RMS norm + if SGL_RMS_NORM is not None: + _ = SGL_RMS_NORM(allreduce_out, rms_gamma, rms_eps) + else: + rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) + rms.weight.data = rms_gamma + _ = rms.forward_native(allreduce_out) + + +def standard_allreduce_rmsnorm_fp8_quant( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rms_gamma: torch.Tensor, + rms_eps: float, + scale_factor: torch.Tensor, + norm_out: Optional[torch.Tensor] = None, + quant_out: Optional[torch.Tensor] = None, +): + """Standard allreduce + rmsnorm + FP8 quantization.""" + # All-reduce first + allreduce_out = tensor_model_parallel_all_reduce(input_tensor) + + # Then RMS norm + static FP8 quantization + if residual is not None: + if SGL_FUSED_ADD_RMS_NORM is not None: + SGL_FUSED_ADD_RMS_NORM(allreduce_out, residual, rms_gamma, rms_eps) + quant_out, _ = static_quant_fp8( + allreduce_out, scale_factor, repeat_scale=False + ) + else: + rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) + rms.weight.data = rms_gamma + normed, _ = rms.forward_native(allreduce_out, residual) + quant_out, _ = static_quant_fp8(normed, scale_factor, repeat_scale=False) + return quant_out, residual + else: + if SGL_RMS_NORM is not None: + normed = SGL_RMS_NORM(allreduce_out, rms_gamma, rms_eps) + else: + rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) + rms.weight.data = rms_gamma + normed = rms.forward_native(allreduce_out) + quant_out, _ = static_quant_fp8(normed, scale_factor, repeat_scale=False) + return quant_out + + +def standard_allreduce_rmsnorm_fp4_quant( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rms_gamma: torch.Tensor, + rms_eps: float, + input_global_scale: torch.Tensor, + quant_out: torch.Tensor, + output_scale: torch.Tensor, + norm_out: Optional[torch.Tensor] = None, +): + """Standard allreduce + rmsnorm + FP4 quantization.""" + + # All-reduce first + allreduce_out = tensor_model_parallel_all_reduce(input_tensor) + + # Then RMS norm + if residual is not None: + if SGL_FUSED_ADD_RMS_NORM is not None: + SGL_FUSED_ADD_RMS_NORM(allreduce_out, residual, rms_gamma, rms_eps) + quant_input = allreduce_out + else: + rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) + rms.weight.data = rms_gamma + quant_input, _ = rms.forward_native(allreduce_out, residual) + residual_out = residual + else: + if SGL_RMS_NORM is not None: + quant_input = SGL_RMS_NORM(allreduce_out, rms_gamma, rms_eps) + else: + rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) + rms.weight.data = rms_gamma + quant_input = rms.forward_native(allreduce_out) + residual_out = allreduce_out + + # Finally FP4 quantization + if SGL_SCALED_FP4_QUANT is None: + raise RuntimeError("scaled_fp4_quant is not available on this platform") + quant_res, output_scale_res = SGL_SCALED_FP4_QUANT(quant_input, input_global_scale) + if residual is not None: + return quant_res, residual_out, output_scale_res + else: + return quant_res, quant_input + + +def standard_allreduce_rmsnorm_native( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rmsnorm_layer: RMSNorm, + norm_out: Optional[torch.Tensor] = None, +): + """Standard allreduce + rmsnorm operations using native RMSNorm forward.""" + # All-reduce first + allreduce_out = tensor_model_parallel_all_reduce(input_tensor) + # Apply native RMSNorm + if residual is not None: + result = rmsnorm_layer.forward_native(allreduce_out, residual) + return result # Returns (norm_out, residual_out) + else: + result = rmsnorm_layer.forward_native(allreduce_out) + return result # Returns norm_out + + +def standard_allreduce_rmsnorm_fp8_quant_native( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rmsnorm_layer: RMSNorm, + scale_factor: torch.Tensor, + norm_out: Optional[torch.Tensor] = None, + quant_out: Optional[torch.Tensor] = None, +): + """Standard allreduce + rmsnorm + FP8 quantization using native implementations.""" + # All-reduce first + allreduce_out = tensor_model_parallel_all_reduce(input_tensor) + + # Apply native RMSNorm + if residual is not None: + norm_out, residual_out = rmsnorm_layer.forward_native(allreduce_out, residual) + else: + norm_out = rmsnorm_layer.forward_native(allreduce_out) + residual_out = allreduce_out + + # Apply native FP8 quantization + quant_out, _ = static_quant_fp8(norm_out, scale_factor, repeat_scale=False) + + if residual is not None: + return quant_out, residual_out + else: + return quant_out + + +def standard_allreduce_rmsnorm_fp4_quant_native( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rmsnorm_layer: RMSNorm, + input_global_scale: torch.Tensor, + quant_out: torch.Tensor, + output_scale: torch.Tensor, + norm_out: Optional[torch.Tensor] = None, +): + """Standard allreduce + rmsnorm + FP4 quantization using native RMSNorm.""" + # All-reduce first + allreduce_out = tensor_model_parallel_all_reduce(input_tensor) + + # Apply native RMSNorm + if residual is not None: + norm_out, residual_out = rmsnorm_layer.forward_native(allreduce_out, residual) + quant_input = norm_out + else: + norm_out = rmsnorm_layer.forward_native(allreduce_out) + quant_input = norm_out + residual_out = allreduce_out + + # Apply FP4 quantization (still using fused CUDA op as there's no native FP4) + if SGL_SCALED_FP4_QUANT is None: + raise RuntimeError("scaled_fp4_quant is not available on this platform") + quant_res, output_scale_res = SGL_SCALED_FP4_QUANT(quant_input, input_global_scale) + + if residual is not None: + return quant_res, residual_out, output_scale_res + else: + return quant_res, norm_out + + +# Compiled versions of native functions +@torch.compile +def standard_allreduce_rmsnorm_native_compiled( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rmsnorm_layer: RMSNorm, + norm_out: Optional[torch.Tensor] = None, +): + """Compiled version of standard allreduce + rmsnorm.""" + return standard_allreduce_rmsnorm_native( + input_tensor, residual, rmsnorm_layer, norm_out + ) + + +@torch.compile +def standard_allreduce_rmsnorm_fp8_quant_native_compiled( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rmsnorm_layer: RMSNorm, + scale_factor: torch.Tensor, + norm_out: Optional[torch.Tensor] = None, + quant_out: Optional[torch.Tensor] = None, +): + """Compiled version of standard allreduce + rmsnorm + FP8 quantization.""" + return standard_allreduce_rmsnorm_fp8_quant_native( + input_tensor, + residual, + rmsnorm_layer, + scale_factor, + norm_out, + quant_out, + ) + + +@torch.compile +def standard_allreduce_rmsnorm_fp4_quant_native_compiled( + input_tensor: torch.Tensor, + residual: Optional[torch.Tensor], + rmsnorm_layer: RMSNorm, + input_global_scale: torch.Tensor, + quant_out: torch.Tensor, + output_scale: torch.Tensor, + norm_out: Optional[torch.Tensor] = None, +): + """Compiled version of standard allreduce + rmsnorm + FP4 quantization.""" + return standard_allreduce_rmsnorm_fp4_quant_native( + input_tensor, + residual, + rmsnorm_layer, + input_global_scale, + quant_out, + output_scale, + norm_out, + ) + + +def create_test_tensors( + seq_len: int, hidden_dim: int, dtype: torch.dtype, use_residual: bool = True +): + """Create test tensors for benchmarking.""" + input_tensor = torch.randn(seq_len, hidden_dim, dtype=dtype) + residual = ( + torch.randn_like(input_tensor) + if use_residual + else torch.zeros_like(input_tensor) + ) + rms_gamma = torch.ones(hidden_dim, dtype=dtype) + norm_out = None if use_residual else torch.empty_like(input_tensor) + + # Quantization scales + scale_fp8 = torch.tensor(1.0, dtype=torch.float32) + scale_fp4 = torch.tensor(1.0, dtype=torch.float32) + quant_out_fp8 = torch.empty_like(input_tensor, dtype=FP8_DTYPE) + # Pre-allocate FP4 output tensors (to avoid allocation overhead in benchmarks) + fp4_quant_out = torch.empty((seq_len, hidden_dim // 2), dtype=torch.uint8) + fp4_output_scale = torch.empty((128, 4), dtype=torch.int32) + + return ( + input_tensor, + norm_out, + residual, + rms_gamma, + scale_fp8, + quant_out_fp8, + scale_fp4, + fp4_quant_out, + fp4_output_scale, + ) + + +def benchmark_operation( + operation_func, *args, warmup: int = 5, trials: int = 20, **kwargs +): + """Benchmark a single operation using CUDA graphs.""" + # Warmup before graph capture + for _ in range(warmup): + operation_func(*args, **kwargs) + torch.cuda.synchronize() + + # Create CUDA graph + graph = torch.cuda.CUDAGraph() + num_op_per_cudagraph = 10 + + # Use sglang's graph_capture to make tensor_model_parallel_all_reduce graph-safe + with graph_capture() as graph_capture_context: + with torch.cuda.graph(graph, stream=graph_capture_context.stream): + for _ in range(num_op_per_cudagraph): + operation_func(*args, **kwargs) + + # Graph warmup + torch.cuda.synchronize() + for _ in range(warmup): + graph.replay() + + # Benchmark with CUDA graph + torch.cuda.synchronize() + start_time = time.perf_counter() + + for _ in range(trials // num_op_per_cudagraph): + # operation_func(*args, **kwargs) + graph.replay() + + torch.cuda.synchronize() + end_time = time.perf_counter() + + avg_time_ms = ((end_time - start_time) / trials) * 1000 + return avg_time_ms + + +def run_benchmarks( + seq_len: int, + hidden_dim: int, + dtype: torch.dtype, + use_residual: bool, + allreduce_params: Optional[FlashInferFusedAllReduceParams], + quant_mode: str = "all", + disable_oneshot: bool = False, +): + """Run all benchmarks for given configuration. + + Args: + quant_mode: "none", "fp8_only", "fp4_only", or "all" + """ + ( + input_tensor, + norm_out, + residual, + rms_gamma, + scale_fp8, + quant_out_fp8, + scale_fp4, + fp4_quant_out, + fp4_output_scale, + ) = create_test_tensors(seq_len, hidden_dim, dtype, use_residual) + + rms_eps = 1e-6 + results = {} + + # Create RMSNorm once for native benchmarks + rmsnorm_layer = RMSNorm(hidden_dim, eps=rms_eps) + rmsnorm_layer.weight.data = rms_gamma + + if quant_mode in ["all", "none"]: + # Standard AllReduce + RMSNorm + try: + time_ms = benchmark_operation( + standard_allreduce_rmsnorm, + input_tensor, + norm_out=norm_out, + residual=residual, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + ) + results["standard_allreduce_rmsnorm"] = time_ms + except Exception as e: + logger.error("Standard AllReduce+RMSNorm failed: %s", e) + results["standard_allreduce_rmsnorm"] = float("inf") + + # Standard AllReduce + RMSNorm Native Compiled + try: + time_ms = benchmark_operation( + standard_allreduce_rmsnorm_native_compiled, + input_tensor, + residual=residual, + rmsnorm_layer=rmsnorm_layer, + norm_out=norm_out, + ) + results["standard_allreduce_rmsnorm_native_compiled"] = time_ms + except Exception as e: + logger.error("Standard AllReduce+RMSNorm Native Compiled failed: %s", e) + results["standard_allreduce_rmsnorm_native_compiled"] = float("inf") + + # FlashInfer Fused AllReduce + RMSNorm Oneshot + if flashinfer_comm is not None and allreduce_params is not None: + try: + if not disable_oneshot: + time_ms = benchmark_operation( + flashinfer_fused_allreduce_rmsnorm, + input_tensor, + residual=residual, + norm_out=norm_out, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + allreduce_params=allreduce_params, + use_oneshot=True, + ) + results["flashinfer_fused_allreduce_rmsnorm_oneshot"] = time_ms + except Exception as e: + logger.error("FlashInfer Fused AllReduce+RMSNorm Oneshot failed: %s", e) + results["flashinfer_fused_allreduce_rmsnorm_oneshot"] = float("inf") + + # FlashInfer Fused AllReduce + RMSNorm Two-shot + try: + time_ms = benchmark_operation( + flashinfer_fused_allreduce_rmsnorm, + input_tensor, + residual=residual, + norm_out=norm_out, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + allreduce_params=allreduce_params, + use_oneshot=False, + ) + results["flashinfer_fused_allreduce_rmsnorm_twoshot"] = time_ms + except Exception as e: + logger.error( + "FlashInfer Fused AllReduce+RMSNorm Two-shot failed: %s", e + ) + results["flashinfer_fused_allreduce_rmsnorm_twoshot"] = float("inf") + + if quant_mode in ["all", "fp8_only"]: + # Standard AllReduce + RMSNorm + FP8 Quant + try: + time_ms = benchmark_operation( + standard_allreduce_rmsnorm_fp8_quant, + input_tensor, + norm_out=norm_out, + residual=residual, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + scale_factor=scale_fp8, + quant_out=quant_out_fp8, + ) + results["standard_allreduce_rmsnorm_fp8_quant"] = time_ms + except Exception as e: + logger.error("Standard AllReduce+RMSNorm+FP8 failed: %s", e) + results["standard_allreduce_rmsnorm_fp8_quant"] = float("inf") + + # Standard AllReduce + RMSNorm + FP8 Quant Native Compiled + try: + time_ms = benchmark_operation( + standard_allreduce_rmsnorm_fp8_quant_native_compiled, + input_tensor, + residual=residual, + rmsnorm_layer=rmsnorm_layer, + # quant_fp8_layer removed in sglang version; static_quant_fp8 is used within the function + scale_factor=scale_fp8, + norm_out=norm_out, + quant_out=quant_out_fp8, + ) + results["standard_allreduce_rmsnorm_fp8_quant_native_compiled"] = time_ms + except Exception as e: + logger.error("Standard AllReduce+RMSNorm+FP8 Native Compiled failed: %s", e) + results["standard_allreduce_rmsnorm_fp8_quant_native_compiled"] = float( + "inf" + ) + + # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Oneshot + if flashinfer_comm is not None and allreduce_params is not None: + try: + if not disable_oneshot: + time_ms = benchmark_operation( + flashinfer_fused_allreduce_rmsnorm_fp8_quant, + input_tensor, + norm_out=norm_out, + residual=residual, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + scale_factor=scale_fp8, + quant_out=quant_out_fp8, + allreduce_params=allreduce_params, + use_oneshot=True, + ) + results["flashinfer_fused_allreduce_rmsnorm_fp8_quant_oneshot"] = ( + time_ms + ) + except Exception as e: + logger.error( + "FlashInfer Fused AllReduce+RMSNorm+FP8 Oneshot failed: %s", + e, + ) + results["flashinfer_fused_allreduce_rmsnorm_fp8_quant_oneshot"] = float( + "inf" + ) + # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Two-shot + try: + time_ms = benchmark_operation( + flashinfer_fused_allreduce_rmsnorm_fp8_quant, + input_tensor, + norm_out=norm_out, + residual=residual, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + scale_factor=scale_fp8, + quant_out=quant_out_fp8, + allreduce_params=allreduce_params, + use_oneshot=False, + ) + results["flashinfer_fused_allreduce_rmsnorm_fp8_quant_twoshot"] = ( + time_ms + ) + except Exception as e: + logger.error( + "FlashInfer Fused AllReduce+RMSNorm+FP8 Two-shot failed: %s", + e, + ) + results["flashinfer_fused_allreduce_rmsnorm_fp8_quant_twoshot"] = float( + "inf" + ) + + if quant_mode in ["all", "fp4_only"]: + # Standard AllReduce + RMSNorm + FP4 Quant + try: + time_ms = benchmark_operation( + standard_allreduce_rmsnorm_fp4_quant, + input_tensor, + norm_out=norm_out, + residual=residual, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + input_global_scale=scale_fp4, + quant_out=fp4_quant_out, + output_scale=fp4_output_scale, + ) + results["standard_allreduce_rmsnorm_fp4_quant"] = time_ms + except Exception as e: + logger.error("Standard AllReduce+RMSNorm+FP4 failed: %s", e) + results["standard_allreduce_rmsnorm_fp4_quant"] = float("inf") + + # Standard AllReduce + RMSNorm + FP4 Quant Native Compiled + try: + time_ms = benchmark_operation( + standard_allreduce_rmsnorm_fp4_quant_native_compiled, + input_tensor, + residual=residual, + rmsnorm_layer=rmsnorm_layer, + input_global_scale=scale_fp4, + quant_out=fp4_quant_out, + output_scale=fp4_output_scale, + norm_out=norm_out, + ) + results["standard_allreduce_rmsnorm_fp4_quant_native_compiled"] = time_ms + except Exception as e: + logger.error("Standard AllReduce+RMSNorm+FP4 Native Compiled failed: %s", e) + results["standard_allreduce_rmsnorm_fp4_quant_native_compiled"] = float( + "inf" + ) + + # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Oneshot + if flashinfer_comm is not None and allreduce_params is not None: + try: + if not disable_oneshot: + time_ms = benchmark_operation( + flashinfer_fused_allreduce_rmsnorm_fp4_quant, + input_tensor, + residual=residual, + norm_out=norm_out, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + input_global_scale=scale_fp4, + allreduce_params=allreduce_params, + quant_out=fp4_quant_out, + output_scale=fp4_output_scale, + use_oneshot=True, + ) + results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_oneshot"] = ( + time_ms + ) + except Exception as e: + logger.error( + "FlashInfer Fused AllReduce+RMSNorm+FP4 Oneshot failed: %s", + e, + ) + results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_oneshot"] = float( + "inf" + ) + + # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Two-shot + if flashinfer_comm is not None and allreduce_params is not None: + try: + time_ms = benchmark_operation( + flashinfer_fused_allreduce_rmsnorm_fp4_quant, + input_tensor, + residual=residual, + norm_out=norm_out, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + input_global_scale=scale_fp4, + allreduce_params=allreduce_params, + quant_out=fp4_quant_out, + output_scale=fp4_output_scale, + use_oneshot=False, + ) + results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = ( + time_ms + ) + except Exception as e: + logger.error( + "FlashInfer Fused AllReduce+RMSNorm+FP4 Two-shot failed: %s", + e, + ) + results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = float( + "inf" + ) + + return results + + +def prepare_results_with_speedups(results_dict): + """Prepare results with speedup calculations based on dynamic baseline selection.""" + prepared_results = [] + + # Determine the fastest baseline for each operation type + def get_fastest_baseline(op_name, results_dict): + """Get the fastest baseline between standard and native_compiled versions.""" + if "fp8_quant" in op_name: + candidates = [ + "standard_allreduce_rmsnorm_fp8_quant", + "standard_allreduce_rmsnorm_fp8_quant_native_compiled", + ] + elif "fp4_quant" in op_name: + candidates = [ + "standard_allreduce_rmsnorm_fp4_quant", + "standard_allreduce_rmsnorm_fp4_quant_native_compiled", + ] + else: + candidates = [ + "standard_allreduce_rmsnorm", + "standard_allreduce_rmsnorm_native_compiled", + ] + + # Find the fastest among available candidates + fastest_time = float("inf") + fastest_baseline = None + + for candidate in candidates: + if ( + candidate in results_dict + and results_dict[candidate] != float("inf") + and results_dict[candidate] < fastest_time + ): + fastest_time = results_dict[candidate] + fastest_baseline = candidate + + return fastest_baseline + + # Create dynamic baseline mapping + dynamic_baseline_mapping = {} + for op_name in results_dict: + if ( + op_name.startswith("flashinfer_") + or op_name.startswith("standard_") + and not op_name.endswith("_native_compiled") + ): + dynamic_baseline_mapping[op_name] = get_fastest_baseline( + op_name, results_dict + ) + + for op_name, time_ms in results_dict.items(): + if time_ms == float("inf"): + speedup_str = "FAILED" + time_str = "FAILED" + else: + time_str = f"{time_ms:.3f}" + # Find the appropriate baseline for this operation + baseline_op = dynamic_baseline_mapping.get(op_name) + if baseline_op and baseline_op in results_dict: + baseline_time = results_dict[baseline_op] + if baseline_time != float("inf") and baseline_time > 0: + speedup = baseline_time / time_ms + speedup_str = f"{speedup:.2f}x" + else: + speedup_str = "N/A" + else: + # For baseline operations, determine if this is the fastest baseline + if op_name.endswith("_native_compiled") or ( + op_name.startswith("standard_") + and not op_name.endswith("_native_compiled") + ): + fastest_baseline = get_fastest_baseline(op_name, results_dict) + if fastest_baseline == op_name: + speedup_str = "baseline" + else: + if fastest_baseline and fastest_baseline in results_dict: + baseline_time = results_dict[fastest_baseline] + if baseline_time != float("inf") and baseline_time > 0: + speedup = baseline_time / time_ms + speedup_str = f"{speedup:.2f}x" + else: + speedup_str = "N/A" + else: + speedup_str = "N/A" + else: + speedup_str = "N/A" + + prepared_results.append( + { + "operation": op_name, + "time_ms": time_ms, + "time_str": time_str, + "speedup_str": speedup_str, + } + ) + + return prepared_results + + +def print_results(results_dict, seq_len, hidden_dim, dtype, use_residual, quant_mode): + """Print benchmark results in a formatted table.""" + print(f"\n{'=' * 80}") + print(f"Results: seq_len={seq_len}, hidden_dim={hidden_dim}") + print( + f"dtype={dtype}, residual={'yes' if use_residual else 'no'}, " + f"quant_mode={quant_mode}" + ) + print(f"{'=' * 80}") + print(f"{'Operation':<50} {'Time (ms)':<12} {'Speedup':<10}") + print(f"{'-' * 80}") + + # Prepare results with speedup calculations + prepared_results = prepare_results_with_speedups(results_dict) + + for result in prepared_results: + if result["time_ms"] == float("inf"): + time_display = result["time_str"] + else: + time_display = f"{result['time_ms']:.3f}" + + print( + f"{result['operation']:<50} {time_display:<12} {result['speedup_str']:<10}" + ) + + +def format_results_markdown( + all_results: list[dict], world_size: int, args: argparse.Namespace +) -> str: + """Format all benchmark results as markdown.""" + markdown = f"""# FlashInfer Fused Collective Operations Benchmark Results + +**World Size:** {world_size} +**Hidden Dimension:** {args.hidden_dim} +**Warmup Iterations:** {args.warmup} +**Benchmark Trials:** {args.trials} +**Quantization Mode:** {all_results[0]["quant_mode"] if all_results else "N/A"} + +--- + +""" + + for result in all_results: + seq_len = result["seq_len"] + dtype = result["dtype"] + use_residual = result["use_residual"] + results_dict = result["results"] + + residual_str = "with residual" if use_residual else "no residual" + + markdown += f""" +## Configuration: seq_len={seq_len}, dtype={dtype}, {residual_str} + +| Operation | Time (ms) | Speedup | +|-----------|-----------|---------| +""" + + # Prepare results with speedup calculations + prepared_results = prepare_results_with_speedups(results_dict) + + for result in prepared_results: + # Format operation name for better readability + formatted_op_name = result["operation"].replace("_", " ").title() + markdown += f"| {formatted_op_name} | {result['time_str']} |" + markdown += f"{result['speedup_str']} |\n" + + markdown += "\n" + + return markdown + + +def save_results_to_file( + all_results: list[dict], world_size: int, args: argparse.Namespace, rank: int +): + """Save benchmark results to markdown file (only on rank 0).""" + if rank != 0: + return + + if not all_results: + logger.warning("No results to save") + return + + output_path = args.output_file + + try: + markdown_content = format_results_markdown(all_results, world_size, args) + + with open(output_path, "w") as f: + f.write(markdown_content) + + except Exception as e: + logger.error("Failed to save results to file: %s", e) + + +def main(): + parser = argparse.ArgumentParser( + description="Benchmark fused collective operations" + ) + parser.add_argument( + "--seq-lens", + type=int, + nargs="+", + default=[128, 512, 1024, 2048], + help="Sequence lengths to test", + ) + parser.add_argument( + "--hidden-dim", type=int, default=8192, help="Hidden dimension size" + ) + parser.add_argument( + "--dtypes", + type=str, + nargs="+", + default=["bfloat16"], + choices=["float16", "bfloat16", "float32"], + help="Data types to test", + ) + parser.add_argument( + "--no-residual", + action="store_true", + help="Skip residual connection tests", + ) + + # Quantization mode options (mutually exclusive with --no-quant) + quant_group = parser.add_mutually_exclusive_group() + quant_group.add_argument( + "--no-quant", action="store_true", help="Skip all quantization tests" + ) + quant_group.add_argument( + "--quant-fp8", action="store_true", help="Only run FP8 quantization tests" + ) + quant_group.add_argument( + "--quant-fp4", action="store_true", help="Only run FP4 quantization tests" + ) + quant_group.add_argument( + "--quant-all", + action="store_true", + help="Run all quantization tests (default)", + ) + + parser.add_argument( + "--disable-oneshot", + action="store_true", + help="Disable oneshot mode for FlashInfer operations", + ) + parser.add_argument( + "--warmup", type=int, default=5, help="Number of warmup iterations" + ) + parser.add_argument( + "--trials", type=int, default=20, help="Number of benchmark trials" + ) + parser.add_argument( + "--output-file", + type=str, + help="""Output file path for markdown results + (default: benchmark_results_.md) + """, + ) + + args = parser.parse_args() + + # Check if running with torchrun (required for collective operations) + if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ: + raise RuntimeError( + "Must run with torchrun for distributed benchmarking. " + "Example: torchrun --nproc_per_node=2 benchmark_fused_collective.py" + ) + + # Initialize distributed environment + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + torch.set_default_device(device) + + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=rank, + backend="nccl", + ) + initialize_model_parallel(tensor_model_parallel_size=world_size) + + # Validate world size (must be > 1 for collective operations) + if world_size <= 1: + raise ValueError( + "World size must be > 1 for collective operations benchmarking. " + f"Current world size: {world_size}. Use torchrun with --nproc_per_node > 1." + ) + + # Determine quantization mode + if args.no_quant: + quant_mode = "none" + elif args.quant_fp8: + quant_mode = "fp8_only" + elif args.quant_fp4: + quant_mode = "fp4_only" + else: # args.quant_all or default + quant_mode = "all" + + if rank == 0: + logger.info("Running benchmark with world_size=%s, rank=%s", world_size, rank) + logger.info("Quantization mode: %s", quant_mode) + if flashinfer_comm is not None: + oneshot_status = "enabled" if not args.disable_oneshot else "disabled" + logger.info( + "FlashInfer available - will benchmark fused operations (oneshot: %s)", + oneshot_status, + ) + else: + logger.info( + "FlashInfer not available - only benchmarking standard operations" + ) + + # Convert dtype strings to torch dtypes + dtype_map = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, + } + dtypes = [dtype_map[dt] for dt in args.dtypes] + + # Test configurations + residual_options = [True] if not args.no_residual else [False] + if not args.no_residual: + residual_options.append(False) + + configs = list(itertools.product(args.seq_lens, dtypes, residual_options)) + + # Setup FlashInfer workspace if available + ipc_handles = None + allreduce_params = None + + if flashinfer_comm is not None: + # Use the largest hidden dimension for workspace setup + max_num_token = _FI_MAX_SIZES.get(world_size) // ( + args.hidden_dim * world_size * 2 + ) + + ipc_handles, workspace_tensor = setup_flashinfer_workspace( + world_size, rank, args.hidden_dim, max_num_token + ) + + if workspace_tensor is not None: + allreduce_params = FlashInferFusedAllReduceParams( + rank=rank, + world_size=world_size, + max_token_num=max_num_token, + ) + + # Collect all results for markdown export + all_results = [] + + try: + # Run benchmarks + for seq_len, dtype, use_residual in configs: + if rank == 0: + logger.info( + "\nTesting: seq_len=%s, hidden_dim=%s, dtype=%s, residual=%s", + seq_len, + args.hidden_dim, + dtype, + use_residual, + ) + + results = run_benchmarks( + seq_len, + args.hidden_dim, + dtype, + use_residual, + allreduce_params, + quant_mode=quant_mode, + disable_oneshot=args.disable_oneshot, + ) + + # Store results for markdown export + if rank == 0: + all_results.append( + { + "seq_len": seq_len, + "hidden_dim": args.hidden_dim, + "dtype": str(dtype).replace("torch.", ""), + "use_residual": use_residual, + "quant_mode": quant_mode, + "results": results, + } + ) + + print_results( + results, + seq_len, + args.hidden_dim, + dtype, + use_residual, + quant_mode, + ) + + # Save results to markdown file + if args.output_file and rank == 0: + save_results_to_file(all_results, world_size, args, rank) + + finally: + # Cleanup + if ipc_handles is not None: + cleanup_flashinfer_workspace(ipc_handles) + + with contextlib.suppress(Exception): + dist.barrier() + cleanup_dist_env_and_memory(shutdown_ray=False) + + +if __name__ == "__main__": + main() diff --git a/benchmark/kernels/fused_moe_triton/README.md b/benchmark/kernels/fused_moe_triton/README.md new file mode 100644 index 000000000..48598854a --- /dev/null +++ b/benchmark/kernels/fused_moe_triton/README.md @@ -0,0 +1,76 @@ +## Tuning Triton MoE Kernels + +This directory contains benchmarking tools for MoE (Mixture of Experts) kernels. + +### Tuning Tool + +- `tuning_fused_moe_triton.py`: A tool for tuning the `fused_moe_triton` kernel. Adapted from [vllm's benchmark_moe.py](https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py), with added support for various model architectures. + +Example usage: +```bash +# Tune Mixtral-8x7B with default settings +python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ + --model mistralai/Mixtral-8x7B-Instruct-v0.1 \ + --tune + +# Tune Qwen2-57B with FP8 and TP=4 +python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ + --model Qwen/Qwen2-57B-A14B-Instruct \ + --tp-size 4 \ + --dtype fp8_w8a8 \ + --tune + +# Tune Qwen3-235B-A22B-FP8 and TP=4 +python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ + --model Qwen/Qwen3-235B-A22B-FP8 \ + --tp-size 4 \ + --dtype fp8_w8a8 \ + --tune + +# Tune DeepSeek-V3 with FP8 and TP=8 +python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ + --model deepseek-ai/DeepSeek-V3-0324 \ + --tp-size 8 \ + --dtype fp8_w8a8 \ + --tune + +# Tune DeepSeek-R1 with channel-wise INT8 and TP=16 +python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ + --model meituan/DeepSeek-R1-Channel-INT8 \ + --tp-size 16 \ + --dtype int8_w8a8 \ + --tune +``` + +After tuning, a configuration file (e.g., `E=64,N=640,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json`) will be generated in the current directory. You can move this file to `sglang/srt/layers/fused_moe_triton/configs/triton_version` dir to use it in `sglang`. + +### Performance Comparison Tool + +- `benchmark_vllm_vs_sglang_fused_moe_triton.py`: A tool for comparing the performance of fused MoE kernels between vllm and sglang implementations. Supports various model architectures and data types. + +Example usage: +```bash +# Compare with default settings (Mixtral model) +python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py + +# Compare with FP8 mode for Qwen2-57B +python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py \ + --model Qwen/Qwen2-57B-A14B-Instruct \ + --use-fp8-w8a8 + +# Compare with custom TP size +python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py \ + --model deepseek-ai/DeepSeek-V3-0324 \ + --tp-size 8 + +# Compare with custom TP size +python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py \ + --model deepseek-ai/DeepSeek-V3-0324 \ + --tp-size 8 +``` + +The benchmark results will be saved as plots and data files in the specified output directory (default: `./configs/benchmark_ops/vllm_sglang_fused_moe/`). + +- `benchmark_torch_compile_fused_moe.py`: A tool for benchmarking the performance of the fused MoE kernel with `torch.compile` and original fused MoE kernel. + +Usage is the same as `benchmark_vllm_vs_sglang_fused_moe_triton.py`, note that `torch.compile` does not support `fp8_w8a8` and `int8_w8a8` fused_moe_kernel. diff --git a/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py b/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py new file mode 100644 index 000000000..7621628c1 --- /dev/null +++ b/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py @@ -0,0 +1,292 @@ +# python3 benchmark/kernels/fused_moe_triton/sglang_fused_moe_triton.py --model /DeepSeek-V3/ --tp-size 8 +import argparse + +import torch +import triton +from transformers import AutoConfig + +from sglang.srt.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, +) +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + fused_moe as fused_moe_sglang, +) +from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import ( + triton_kernel_moe_forward, +) +from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig +from sglang.srt.layers.moe.topk import TopK, TopKConfig, select_experts + + +def get_model_config(model_name: str, tp_size: int): + """Get model configuration parameters""" + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + + if config.architectures[0] == "Qwen2MoeForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "Qwen3MoeForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] in [ + "DeepseekV2ForCausalLM", + "DeepseekV3ForCausalLM", + "Glm4MoeForCausalLM", + ]: + E = ( + config.n_routed_experts + 1 + if config.architectures[0] in ["DeepseekV3ForCausalLM"] + else config.n_routed_experts + ) + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + else: + # Default: Mixtral + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + + block_shape = None + if ( + hasattr(config, "quantization_config") + and "weight_block_size" in config.quantization_config + ): + block_shape = config.quantization_config["weight_block_size"] + assert len(block_shape) == 2 + + shape_configs = { + "num_experts": E, + "topk": topk, + "hidden_size": config.hidden_size, + "shard_intermediate_size": shard_intermediate_size, + "dtype": config.torch_dtype, + "block_shape": block_shape, + } + print(f"{shape_configs=}") + return shape_configs + + +def fused_moe_triton_api( + x, + w1, + w2, + input_gating, + topk, +): + topk_op = TopK( + top_k=topk, + renormalize=False, + use_grouped_topk=False, + ) + topk_op.use_triton_kernels = True + triton_topk_output = topk_op.forward_cuda( + hidden_states=x, + router_logits=input_gating, + ) + + moe_runner_config = MoeRunnerConfig( + inplace=False, + ) + return triton_kernel_moe_forward( + x, + w1, + w2, + triton_topk_output, + moe_runner_config, + ) + + +def fused_moe_sglang_api( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, + block_shape=None, +): + topk_output = select_experts( + hidden_states=x, + router_logits=input_gating, + topk_config=TopKConfig(top_k=topk, renormalize=False), + ) + return fused_moe_sglang( + x, + w1, + w2, + topk_output, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=list([128, 256, 512, 1024, 2048, 4096, 8192]), + line_arg="provider", + line_vals=[ + "sglang_fused_moe_triton_v340", + "sglang_fused_moe_triton", + ], + line_names=[ + "sglang_fused_moe_triton_v340", + "sglang_fused_moe_triton", + ], + styles=[ + ("blue", "-"), + ("green", "-"), + ], + ylabel="Time (ms)", + plot_name="fused-moe-performance", + args={}, + ) +) +def benchmark( + batch_size, + provider, + model_config, + use_fp8_w8a8=False, + use_cuda_graph: bool = False, +): + print(f"benchmark {provider} with batch_size={batch_size}") + torch.set_default_device("cuda") + torch.cuda.manual_seed_all(0) + + num_tokens = batch_size + num_experts = model_config["num_experts"] + hidden_size = model_config["hidden_size"] + shard_intermediate_size = model_config["shard_intermediate_size"] + topk = model_config["topk"] + dtype = model_config["dtype"] + block_shape = model_config["block_shape"] + + x = torch.randn(num_tokens, hidden_size, dtype=dtype) + + w1 = torch.randn(num_experts, shard_intermediate_size, hidden_size, dtype=dtype) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=dtype + ) + + w1_tri = w1.clone() + w2_tri = w2.clone() + w1_tri = w1_tri.transpose(-2, -1).contiguous() + w2_tri = w2_tri.transpose(-2, -1).contiguous() + + input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) + + if provider == "sglang_fused_moe_triton_v340": + api_func = fused_moe_triton_api + api_kwargs = { + "x": x, + "w1": w1_tri, + "w2": w2_tri, + "input_gating": input_gating, + "topk": topk, + } + else: + api_func = fused_moe_sglang_api + api_kwargs = { + "x": x, + "w1": w1, + "w2": w2, + "input_gating": input_gating, + "topk": topk, + "use_fp8_w8a8": use_fp8_w8a8, + "block_shape": block_shape, + } + + # Warmup + for _ in range(10): + _ = api_func(**api_kwargs) + torch.cuda.synchronize() + + if use_cuda_graph: + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + api_func(**api_kwargs) + torch.cuda.synchronize() + + bench_lambda = lambda: graph.replay() + else: + bench_lambda = lambda: api_func(**api_kwargs) + + quantiles = [0.5, 0.2, 0.8] + ms, min_ms, max_ms = triton.testing.do_bench(bench_lambda, quantiles=quantiles) + return ms, min_ms, max_ms + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" + ) + parser.add_argument("--tp-size", type=int, default=2) + parser.add_argument("--use-fp8-w8a8", action="store_true") + parser.add_argument( + "--use-cuda-graph", action="store_true", help="Enable CUDA Graph capture/replay" + ) + parser.add_argument( + "--save-path", + type=str, + default="./configs/benchmark_ops/sglang_fused_moe/", + ) + parser.add_argument("--trust-remote-code", action="store_true") + args = parser.parse_args() + + try: + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group( + backend="nccl" if torch.cuda.is_available() else "gloo", + init_method="tcp://127.0.0.1:23456", + world_size=1, + rank=0, + ) + + init_distributed_environment( + world_size=1, + rank=0, + distributed_init_method="tcp://127.0.0.1:23456", + local_rank=0, + backend="nccl" if torch.cuda.is_available() else "gloo", + ) + + initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + + model_config = get_model_config(args.model, args.tp_size) + benchmark.run( + show_plots=True, + print_data=True, + save_path=args.save_path, + model_config=model_config, + use_fp8_w8a8=args.use_fp8_w8a8, + use_cuda_graph=args.use_cuda_graph, + ) + finally: + destroy_model_parallel() + destroy_distributed_environment() + + +if __name__ == "__main__": + main() diff --git a/benchmark/kernels/fused_moe_triton/benchmark_sum_scale.py b/benchmark/kernels/fused_moe_triton/benchmark_sum_scale.py new file mode 100644 index 000000000..979d2bbd1 --- /dev/null +++ b/benchmark/kernels/fused_moe_triton/benchmark_sum_scale.py @@ -0,0 +1,202 @@ +import torch +import triton +import triton.language as tl +from triton.testing import do_bench + + +@triton.jit +def _moe_sum_reduce_kernel( + input_ptr, + input_stride_0, + input_stride_1, + input_stride_2, + output_ptr, + output_stride_0, + output_stride_1, + token_num: int, + topk_num: int, + hidden_dim: int, + routed_scaling_factor: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DIM: tl.constexpr, + NUM_STAGE: tl.constexpr, +): + input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64) + input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64) + output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64) + + token_block_id = tl.program_id(0) + dim_block_id = tl.program_id(1) + + offs_token = token_block_id * BLOCK_M + tl.arange(0, BLOCK_M) + offs_dim = dim_block_id * BLOCK_DIM + tl.arange(0, BLOCK_DIM) + + mask_token = offs_token < token_num + mask_dim = offs_dim < hidden_dim + + base_ptrs = input_ptr + offs_token[:, None] * input_stride_0 + offs_dim[None, :] + + accumulator = tl.zeros((BLOCK_M, BLOCK_DIM), dtype=tl.float32) + + for i in tl.range(0, topk_num, num_stages=NUM_STAGE): + tile = tl.load( + base_ptrs + i * input_stride_1, + mask=mask_token[:, None] & mask_dim[None, :], + other=0.0, + ) + accumulator += tile.to(tl.float32) + accumulator *= routed_scaling_factor + + # -------- Write back -------- + store_ptrs = output_ptr + offs_token[:, None] * output_stride_0 + offs_dim[None, :] + tl.store( + store_ptrs, + accumulator.to(input_ptr.dtype.element_ty), + mask=mask_token[:, None] & mask_dim[None, :], + ) + + +# _moe_sum_reduce_kernel kernel modified from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/moe_sum_reduce.py +def moe_sum_reduce( + input: torch.Tensor, output: torch.Tensor, routed_scaling_factor: float +): + assert input.is_contiguous() + assert output.is_contiguous() + + token_num, topk_num, hidden_dim = input.shape + assert output.shape[0] == token_num and output.shape[1] == hidden_dim + + BLOCK_M = 1 + BLOCK_DIM = 2048 + NUM_STAGE = 1 + num_warps = 16 + + grid = ( + triton.cdiv(token_num, BLOCK_M), + triton.cdiv(hidden_dim, BLOCK_DIM), + ) + + _moe_sum_reduce_kernel[grid]( + input, + *input.stride(), + output, + *output.stride(), + token_num=token_num, + topk_num=topk_num, + hidden_dim=hidden_dim, + routed_scaling_factor=routed_scaling_factor, + BLOCK_M=BLOCK_M, + BLOCK_DIM=BLOCK_DIM, + NUM_STAGE=NUM_STAGE, + num_warps=num_warps, + ) + return + + +def compute_sum_scaled_baseline( + x: torch.Tensor, out: torch.Tensor, routed_scaling_factor: float +) -> torch.Tensor: + torch.sum(x, dim=1, out=out) + out.mul_(routed_scaling_factor) + return out + + +@torch.compile +def compute_sum_scaled_compiled( + x: torch.Tensor, out: torch.Tensor, routed_scaling_factor: float +) -> torch.Tensor: + torch.sum(x * routed_scaling_factor, dim=1, out=out) + return out + + +def get_benchmark(): + num_tokens_range = [2**i for i in range(0, 13)] + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens"], + x_vals=num_tokens_range, + line_arg="version", + line_vals=["baseline", "compiled", "triton"], + line_names=["Original", "TorchCompile", "TritonKernel"], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="us", + plot_name="sum_scaled_performance", + args={}, + ) + ) + def benchmark(num_tokens, version): + topk = 9 + hidden_size = 4096 + dtype = torch.bfloat16 + scaling_factor = 0.3 + + x = torch.randn(num_tokens, topk, hidden_size, dtype=dtype, device="cuda") + out = torch.empty(num_tokens, hidden_size, dtype=dtype, device="cuda") + + # Warmup + for _ in range(3): + if version == "baseline": + compute_sum_scaled_baseline(x, out, scaling_factor) + elif version == "compiled": + compute_sum_scaled_compiled(x, out, scaling_factor) + else: + moe_sum_reduce(x, out, scaling_factor) + + # Benchmark + quantiles = [0.5, 0.2, 0.8] + if version == "baseline": + ms, min_ms, max_ms = do_bench( + lambda: compute_sum_scaled_baseline(x, out, scaling_factor), + quantiles=quantiles, + ) + elif version == "compiled": + ms, min_ms, max_ms = do_bench( + lambda: compute_sum_scaled_compiled(x, out, scaling_factor), + quantiles=quantiles, + ) + else: + ms, min_ms, max_ms = do_bench( + lambda: moe_sum_reduce(x, out, scaling_factor), quantiles=quantiles + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +def verify_correctness(num_tokens=1024): + x = torch.randn(num_tokens, 9, 4096, device="cuda", dtype=torch.bfloat16) + scaling_factor = 0.3 + + out_baseline = torch.empty_like(x[:, 0]) + compute_sum_scaled_baseline(x, out_baseline, scaling_factor) + + out_compiled = torch.empty_like(out_baseline) + compute_sum_scaled_compiled(x, out_compiled, scaling_factor) + + out_triton = torch.empty_like(out_baseline) + moe_sum_reduce(x, out_triton, scaling_factor) + + if torch.allclose( + out_baseline, out_compiled, atol=1e-2, rtol=1e-2 + ) and torch.allclose(out_baseline, out_triton, atol=1e-2, rtol=1e-2): + print("✅ All implementations match") + else: + print("❌ Implementations differ") + print( + f"Baseline vs Compiled: {(out_baseline - out_compiled).abs().max().item()}" + ) + print(f"Baseline vs Triton: {(out_baseline - out_triton).abs().max().item()}") + + +if __name__ == "__main__": + print("Running correctness verification...") + verify_correctness() + + print("\nRunning performance benchmark...") + benchmark = get_benchmark() + benchmark.run( + print_data=True, + # save_path="./configs/benchmark_ops/sum_scaled/" + ) diff --git a/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py b/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py new file mode 100644 index 000000000..2b4faa24b --- /dev/null +++ b/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py @@ -0,0 +1,305 @@ +# python3 benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py --model /DeepSeek-V3/ --tp-size 8 --use-fp8-w8a8 +import argparse + +import torch +import triton +from torch.nn import functional as F +from transformers import AutoConfig + +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + fused_moe as fused_moe_triton, +) +from sglang.srt.model_executor.cuda_graph_runner import set_torch_compile_config + + +def get_model_config(model_name: str, tp_size: int): + """Get model configuration parameters""" + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + + if config.architectures[0] == "DbrxForCausalLM": + E = config.ffn_config.moe_num_experts + topk = config.ffn_config.moe_top_k + intermediate_size = config.ffn_config.ffn_hidden_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "JambaForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "Qwen2MoeForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "Qwen3MoeForCausalLM": + E = config.n_routed_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]: + E = config.n_routed_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "Llama4ForConditionalGeneration": + E = config.text_config.num_local_experts + topk = config.text_config.num_experts_per_tok + intermediate_size = config.text_config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] in [ + "Grok1ForCausalLM", + "Grok1ImgGen", + "Grok1AForCausalLM", + ]: + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + else: + # Default: Mixtral + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + + shape_configs = { + "num_experts": E, + "topk": topk, + "hidden_size": config.hidden_size, + "shard_intermediate_size": shard_intermediate_size, + "dtype": config.torch_dtype, + } + print(f"{shape_configs=}") + return shape_configs + + +def fused_topk_native( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, +): + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + M, _ = hidden_states.shape + topk_weights = torch.empty( + M, topk, dtype=torch.float32, device=hidden_states.device + ) + topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device) + topk_weights = F.softmax(gating_output.float(), dim=-1) + topk_weights, topk_ids = torch.topk(topk_weights, topk, dim=-1) + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + return topk_weights, topk_ids + + +@torch.compile(dynamic=False) +def fused_moe_torch( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, +) -> torch.Tensor: + assert not use_fp8_w8a8, "Fp8_w8a8 fused_moe is not supported for torch compile" + + topk_weights, topk_ids = fused_topk_native( + hidden_states=x, + gating_output=input_gating, + topk=topk, + renormalize=True, + ) + w13_weights = w1[topk_ids] + w1_weights, w3_weights = torch.chunk(w13_weights, 2, dim=2) + w2_weights = w2[topk_ids] + x1 = torch.einsum("ti,taoi -> tao", x, w1_weights) + x1 = F.silu(x1) + x3 = torch.einsum("ti, taoi -> tao", x, w3_weights) + expert_outs = torch.einsum("tao, taio -> tai", (x1 * x3), w2_weights) + return torch.einsum("tai,ta -> ti", expert_outs, topk_weights.to(expert_outs.dtype)) + + +def fused_moe_torch_compile( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, +): + return fused_moe_torch( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + ) + + +def fused_moe_sglang_api( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, +): + return fused_moe_triton( + x, + w1, + w2, + input_gating, + topk, + renormalize=True, + inplace=True, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + ) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=list(range(1, 5)), + line_arg="provider", + line_vals=[ + "fused_moe_triton", + "fused_moe_torch_compile", + ], + line_names=[ + "fused_moe_triton", + "fused_moe_torch_compile", + ], + styles=[ + ("blue", "-"), + ("green", "-"), + ], + ylabel="Time (ms)", + plot_name="fused-moe-performance", + args={}, + ) +) +def benchmark(batch_size, provider, model_config, use_fp8_w8a8=False): + print(f"benchmark {provider} with batch_size={batch_size}") + torch.set_default_device("cuda") + torch.cuda.manual_seed_all(0) + set_torch_compile_config() + + num_tokens = batch_size + num_experts = model_config["num_experts"] + hidden_size = model_config["hidden_size"] + shard_intermediate_size = model_config["shard_intermediate_size"] + topk = model_config["topk"] + dtype = model_config["dtype"] + + x = torch.randn(num_tokens, hidden_size, dtype=dtype) + + if use_fp8_w8a8: + init_dtype = dtype + w1 = torch.randn( + num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype + ) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype + ) + w1 = w1.to(torch.float8_e4m3fn) + w2 = w2.to(torch.float8_e4m3fn) + w1_scale = torch.randn(num_experts, dtype=torch.float32) + w2_scale = torch.randn(num_experts, dtype=torch.float32) + a1_scale = torch.randn(1, dtype=torch.float32) + a2_scale = torch.randn(1, dtype=torch.float32) + else: + w1 = torch.randn(num_experts, shard_intermediate_size, hidden_size, dtype=dtype) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=dtype + ) + w1_scale = w2_scale = a1_scale = a2_scale = None + + input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) + + # Warmup + api_func = ( + fused_moe_torch_compile + if provider == "fused_moe_torch_compile" + else fused_moe_sglang_api + ) + for _ in range(10): + y = api_func( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + ) + torch.cuda.synchronize() + + quantiles = [0.5, 0.2, 0.8] + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: api_func( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + )[0], + quantiles=quantiles, + ) + return ms, min_ms, max_ms + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" + ) + parser.add_argument("--tp-size", type=int, default=2) + parser.add_argument("--use-fp8-w8a8", action="store_true") + parser.add_argument( + "--save-path", + type=str, + default="./configs/benchmark_ops/fused_moe_torch_compile/", + ) + args = parser.parse_args() + + model_config = get_model_config(args.model, args.tp_size) + benchmark.run( + show_plots=True, + print_data=True, + save_path=args.save_path, + model_config=model_config, + use_fp8_w8a8=args.use_fp8_w8a8, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py b/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py new file mode 100644 index 000000000..6afd7f354 --- /dev/null +++ b/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py @@ -0,0 +1,349 @@ +# python3 benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py --model /DeepSeek-V3/ --tp-size 8 --use-fp8-w8a8 +import argparse + +import torch +import triton +import vllm +from transformers import AutoConfig +from vllm.model_executor.layers.fused_moe.fused_moe import fused_moe as fused_moe_vllm + +from sglang.srt.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, +) +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + fused_moe as fused_moe_sglang, +) + + +def get_model_config(model_name: str, tp_size: int): + """Get model configuration parameters""" + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + + if config.architectures[0] == "DbrxForCausalLM": + E = config.ffn_config.moe_num_experts + topk = config.ffn_config.moe_top_k + intermediate_size = config.ffn_config.ffn_hidden_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "JambaForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "Qwen2MoeForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "Qwen3MoeForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] in [ + "DeepseekV2ForCausalLM", + "DeepseekV3ForCausalLM", + "Glm4MoeForCausalLM", + ]: + E = ( + config.n_routed_experts + 1 + if config.architectures[0] in ["DeepseekV3ForCausalLM"] + else config.n_routed_experts + ) + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "Llama4ForConditionalGeneration": + E = config.text_config.num_local_experts + topk = config.text_config.num_experts_per_tok + intermediate_size = config.text_config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] in [ + "Grok1ForCausalLM", + "Grok1ImgGen", + "Grok1AForCausalLM", + ]: + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + else: + # Default: Mixtral + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + + vllm_version_num = ( + vllm.__version_tuple__[0] * 100 + + vllm.__version_tuple__[1] * 10 + + vllm.__version_tuple__[2] + ) + block_shape = None + if ( + hasattr(config, "quantization_config") + and "weight_block_size" in config.quantization_config + ): + block_shape = config.quantization_config["weight_block_size"] + assert len(block_shape) == 2 + assert ( + vllm_version_num >= 66 + ), "Block-wise quantized fp8 fused_moe is only supported for VLLM>=0.6.6.post1" + + shape_configs = { + "num_experts": E, + "topk": topk, + "hidden_size": config.hidden_size, + "shard_intermediate_size": shard_intermediate_size, + "dtype": config.torch_dtype, + "block_shape": block_shape, + } + print(f"{shape_configs=}") + return shape_configs + + +def fused_moe_vllm_api( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, + block_shape=None, +): + if block_shape is not None: + return fused_moe_vllm( + x, + w1, + w2, + input_gating, + topk, + renormalize=True, + inplace=True, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + else: + return fused_moe_vllm( + x, + w1, + w2, + input_gating, + topk, + renormalize=True, + inplace=True, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + ) + + +def fused_moe_sglang_api( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, + block_shape=None, +): + return fused_moe_sglang( + x, + w1, + w2, + input_gating, + topk, + renormalize=True, + inplace=True, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=list(range(1, 513)), + line_arg="provider", + line_vals=[ + "vllm_fused_moe_triton", + "sglang_fused_moe_triton", + ], + line_names=[ + "vllm_fused_moe_triton", + "sglang_fused_moe_triton", + ], + styles=[ + ("blue", "-"), + ("green", "-"), + ], + ylabel="Time (ms)", + plot_name="fused-moe-performance", + args={}, + ) +) +def benchmark(batch_size, provider, model_config, use_fp8_w8a8=False): + print(f"benchmark {provider} with batch_size={batch_size}") + torch.set_default_device("cuda") + torch.cuda.manual_seed_all(0) + + num_tokens = batch_size + num_experts = model_config["num_experts"] + hidden_size = model_config["hidden_size"] + shard_intermediate_size = model_config["shard_intermediate_size"] + topk = model_config["topk"] + dtype = model_config["dtype"] + block_shape = model_config["block_shape"] + + x = torch.randn(num_tokens, hidden_size, dtype=dtype) + w1_scale = w2_scale = a1_scale = a2_scale = None + + if use_fp8_w8a8: + init_dtype = dtype + w1 = torch.randn( + num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype + ) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype + ) + w1 = w1.to(torch.float8_e4m3fn) + w2 = w2.to(torch.float8_e4m3fn) + + if block_shape is None: + w1_scale = torch.randn(num_experts, dtype=torch.float32) + w2_scale = torch.randn(num_experts, dtype=torch.float32) + a1_scale = torch.randn(1, dtype=torch.float32) + a2_scale = torch.randn(1, dtype=torch.float32) + else: + block_n, block_k = block_shape[0], block_shape[1] + n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n + n_tiles_w2 = (hidden_size + block_n - 1) // block_n + k_tiles_w1 = (hidden_size + block_k - 1) // block_k + k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k + w1_scale = torch.rand( + (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32 + ) + w2_scale = torch.rand( + (num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.float32 + ) + else: + w1 = torch.randn(num_experts, shard_intermediate_size, hidden_size, dtype=dtype) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=dtype + ) + + input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) + + # Warmup + api_func = ( + fused_moe_vllm_api + if provider == "vllm_fused_moe_triton" + else fused_moe_sglang_api + ) + for _ in range(10): + y = api_func( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + torch.cuda.synchronize() + + quantiles = [0.5, 0.2, 0.8] + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: api_func( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + )[0], + quantiles=quantiles, + ) + return ms, min_ms, max_ms + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" + ) + parser.add_argument("--tp-size", type=int, default=2) + parser.add_argument("--use-fp8-w8a8", action="store_true") + parser.add_argument( + "--save-path", + type=str, + default="./configs/benchmark_ops/vllm_sglang_fused_moe/", + ) + args = parser.parse_args() + + try: + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group( + backend="nccl" if torch.cuda.is_available() else "gloo", + init_method="tcp://127.0.0.1:23456", + world_size=1, + rank=0, + ) + + init_distributed_environment( + world_size=1, + rank=0, + distributed_init_method="tcp://127.0.0.1:23456", + local_rank=0, + backend="nccl" if torch.cuda.is_available() else "gloo", + ) + + initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + + model_config = get_model_config(args.model, args.tp_size) + benchmark.run( + show_plots=True, + print_data=True, + save_path=args.save_path, + model_config=model_config, + use_fp8_w8a8=args.use_fp8_w8a8, + ) + finally: + destroy_model_parallel() + destroy_distributed_environment() + + +if __name__ == "__main__": + main() diff --git a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py new file mode 100644 index 000000000..7b52f02a3 --- /dev/null +++ b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py @@ -0,0 +1,599 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py +import argparse +import json +import time +from contextlib import nullcontext +from datetime import datetime +from typing import Any, Dict, List, Tuple, TypedDict + +import ray +import torch +import triton +from ray.experimental.tqdm_ray import tqdm +from transformers import AutoConfig + +from sglang.srt.layers.moe.fused_moe_triton import override_config +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + fused_moe, + get_config_dtype_str, + get_config_file_name, + get_default_config, + get_moe_configs, +) +from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig +from sglang.srt.layers.moe.topk import TopKConfig, select_experts +from sglang.srt.utils import is_hip + +_is_hip = is_hip() + + +class BenchmarkConfig(TypedDict): + BLOCK_SIZE_M: int + BLOCK_SIZE_N: int + BLOCK_SIZE_K: int + GROUP_SIZE_M: int + num_warps: int + num_stages: int + + +def benchmark_config( + config: BenchmarkConfig, + num_tokens: int, + num_experts: int, + shard_intermediate_size: int, + hidden_size: int, + topk: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + block_shape: List[int] = None, + num_iters: int = 100, +) -> float: + init_dtype = torch.float16 if use_fp8_w8a8 else dtype + x = torch.randn(num_tokens, hidden_size, dtype=dtype) + if use_int8_w8a16 or use_int8_w8a8: + w1 = torch.randint( + -127, + 127, + ( + num_experts, + shard_intermediate_size, + hidden_size, + ), + dtype=torch.int8, + ) + w2 = torch.randint( + -127, + 127, + ( + num_experts, + hidden_size, + shard_intermediate_size // 2, + ), + dtype=torch.int8, + ) + else: + w1 = torch.randn( + num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype + ) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype + ) + gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32) + + w1_scale = None + w2_scale = None + a1_scale = None + a2_scale = None + if use_int8_w8a16: + w1_scale = torch.randn( + (num_experts, 2 * shard_intermediate_size), dtype=torch.float32 + ) + w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32) + if use_fp8_w8a8 or use_int8_w8a8: + if use_int8_w8a8 and block_shape is None: + w1_scale = torch.randn( + num_experts, shard_intermediate_size, dtype=torch.float32 + ) + w2_scale = torch.randn(num_experts, hidden_size, dtype=torch.float32) + elif block_shape is None: + w1_scale = torch.randn(num_experts, dtype=torch.float32) + w2_scale = torch.randn(num_experts, dtype=torch.float32) + a1_scale = torch.randn(1, dtype=torch.float32) + a2_scale = torch.randn(1, dtype=torch.float32) + else: + block_n, block_k = block_shape[0], block_shape[1] + n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n + n_tiles_w2 = (hidden_size + block_n - 1) // block_n + k_tiles_w1 = (hidden_size + block_k - 1) // block_k + k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k + w1_scale = torch.rand( + (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32 + ) + w2_scale = torch.rand( + (num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.float32 + ) + + if use_fp8_w8a8: + w1 = w1.to(torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn) + w2 = w2.to(torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn) + + input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) + topk_config = TopKConfig( + top_k=topk, + renormalize=True, + ) + topk_output = select_experts(x, input_gating, topk_config) + + def prepare(i: int): + input_gating = gating_output[i] + new_topk_output = select_experts(x, input_gating, topk_config) + topk_output.topk_weights.copy_(new_topk_output.topk_weights) + topk_output.topk_ids.copy_(new_topk_output.topk_ids) + topk_output.router_logits.copy_(new_topk_output.router_logits) + + def run(): + moe_runner_config = MoeRunnerConfig( + inplace=True, + ) + + with override_config(config): + fused_moe( + x, + w1, + w2, + topk_output, + moe_runner_config=moe_runner_config, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + + # JIT compilation & warmup + run() + torch.cuda.synchronize() + + # Capture 10 invocations with CUDA graph + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(10): + run() + torch.cuda.synchronize() + + # Warmup + for _ in range(5): + graph.replay() + torch.cuda.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + latencies: List[float] = [] + for i in range(num_iters): + prepare(i) + torch.cuda.synchronize() + + start_event.record() + graph.replay() + end_event.record() + end_event.synchronize() + latencies.append(start_event.elapsed_time(end_event)) + avg = sum(latencies) / (num_iters * 10) * 1000 # us + graph.reset() + return avg + + +def get_rocm_configs_compute_bound() -> List[Dict[str, int]]: + configs: List[BenchmarkConfig] = [] + waves_per_eu_range = 0 + for num_stages in [2]: + for block_m in [32, 64, 128, 256]: + for block_k in [32, 64, 128, 256]: + for block_n in [16, 32, 64, 128, 256]: + for num_warps in [1, 2, 4, 8]: + for group_size in [1, 4, 8, 16, 32]: + configs.append( + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_size, + "num_warps": num_warps, + "num_stages": num_stages, + "waves_per_eu": waves_per_eu_range, + } + ) + return configs + + +def get_configs_compute_bound() -> List[Dict[str, int]]: + # Reduced search space for faster tuning. + # TODO(woosuk): Increase the search space and use a performance model to + # prune the search space. + configs: List[BenchmarkConfig] = [] + if _is_hip: + configs = get_rocm_configs_compute_bound() + else: + for num_stages in [2, 3, 4, 5]: + for block_m in [16, 32, 64, 128, 256]: + for block_k in [64, 128, 256]: + for block_n in [32, 64, 128, 256]: + for num_warps in [4, 8]: + for group_size in [1, 16, 32, 64]: + configs.append( + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_size, + "num_warps": num_warps, + "num_stages": num_stages, + } + ) + return configs + + +@ray.remote(num_gpus=1) +class BenchmarkWorker: + + def __init__(self, seed: int) -> None: + torch.set_default_device("cuda") + torch.cuda.manual_seed_all(0) + self.seed = seed + # Get the device ID to allocate tensors and kernels + # on the respective GPU. + self.device_id = int(ray.get_gpu_ids()[0]) + + def benchmark( + self, + num_tokens: int, + num_experts: int, + shard_intermediate_size: int, + hidden_size: int, + topk: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + block_shape: List[int], + ) -> Tuple[Dict[str, int], float]: + torch.cuda.manual_seed_all(0) + dtype_str = get_config_dtype_str( + dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8 + ) + # NOTE(woosuk): The current naming convention uses w2.shape[2], which + # is the intermediate size after silu_and_mul. + block_n = block_shape[0] if block_shape else 0 + block_k = block_shape[1] if block_shape else 0 + op_config = get_moe_configs( + num_experts, shard_intermediate_size // 2, dtype_str, block_n, block_k + ) + if op_config is None: + config = get_default_config( + num_tokens, + num_experts, + shard_intermediate_size, + hidden_size, + topk, + dtype_str, + False, + block_shape, + ) + else: + config = op_config[min(op_config.keys(), key=lambda x: abs(x - num_tokens))] + with torch.cuda.device(self.device_id) if is_hip() else nullcontext(): + kernel_time = benchmark_config( + config, + num_tokens, + num_experts, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + ) + return config, kernel_time + + def tune( + self, + num_tokens: int, + num_experts: int, + shard_intermediate_size: int, + hidden_size: int, + topk: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + block_shape: List[int], + search_space: List[Dict[str, int]], + ) -> Dict[str, int]: + best_config = None + best_time = float("inf") + with torch.cuda.device(self.device_id) if is_hip() else nullcontext(): + for config in tqdm(search_space): + try: + kernel_time = benchmark_config( + config, + num_tokens, + num_experts, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + num_iters=10, + ) + except (triton.runtime.autotuner.OutOfResources, RuntimeError): + # Some configurations may be invalid and fail to compile. + continue + + if kernel_time < best_time: + best_time = kernel_time + best_config = config + now = datetime.now() + print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}") + assert best_config is not None + return best_config + + +def sort_config(config: BenchmarkConfig) -> BenchmarkConfig: + return { + "BLOCK_SIZE_M": config["BLOCK_SIZE_M"], + "BLOCK_SIZE_N": config["BLOCK_SIZE_N"], + "BLOCK_SIZE_K": config["BLOCK_SIZE_K"], + "GROUP_SIZE_M": config["GROUP_SIZE_M"], + "num_warps": config["num_warps"], + "num_stages": config["num_stages"], + **( + {"waves_per_eu": config["waves_per_eu"]} if "waves_per_eu" in config else {} + ), + } + + +def save_configs( + configs: Dict[int, BenchmarkConfig], + num_experts: int, + shard_intermediate_size: int, + hidden_size: int, + topk: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + block_shape: List[int], +) -> None: + dtype_str = get_config_dtype_str( + dtype, + use_int8_w8a16=use_int8_w8a16, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + ) + + # NOTE(woosuk): The current naming convention uses w2.shape[2], which + # is the intermediate size after silu_and_mul. + filename = get_config_file_name( + num_experts, + shard_intermediate_size // 2, + dtype_str, + block_shape, + ) + + print(f"Writing best config to {filename}...") + with open(filename, "w") as f: + json.dump(configs, f, indent=4) + f.write("\n") + + +def main(args: argparse.Namespace): + print(args) + + config = AutoConfig.from_pretrained(args.model, trust_remote_code=True) + if config.architectures[0] == "DbrxForCausalLM": + E = config.ffn_config.moe_num_experts + topk = config.ffn_config.moe_top_k + intermediate_size = config.ffn_config.ffn_hidden_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] == "JambaForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] in ["Qwen2MoeForCausalLM", "Qwen3MoeForCausalLM"]: + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]: + E = ( + config.n_routed_experts + (0 if args.disable_shared_experts_fusion else 1) + if config.architectures[0] in ["DeepseekV3ForCausalLM"] + else config.n_routed_experts + ) + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] == "Llama4ForConditionalGeneration": + E = config.text_config.num_local_experts + ( + 0 if args.disable_shared_experts_fusion else 1 + ) + topk = config.text_config.num_experts_per_tok + intermediate_size = config.text_config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] in [ + "Grok1ForCausalLM", + "Grok1ImgGen", + "Grok1AForCausalLM", + ]: + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] in ["Glm4MoeForCausalLM"]: + E = config.n_routed_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + else: + # Default: Mixtral + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + + hidden_size = getattr(config, "hidden_size", None) or config.text_config.hidden_size + dtype = config.torch_dtype + use_fp8_w8a8 = args.dtype == "fp8_w8a8" + use_int8_w8a8 = args.dtype == "int8_w8a8" + use_int8_w8a16 = args.dtype == "int8_w8a16" + block_shape = None + if ( + hasattr(config, "quantization_config") + and "weight_block_size" in config.quantization_config + ): + block_shape = config.quantization_config["weight_block_size"] + assert len(block_shape) == 2 + + if args.batch_size is None: + batch_sizes = [ + 1, + 2, + 4, + 8, + 16, + 24, + 32, + 48, + 64, + 96, + 128, + 256, + 512, + 1024, + 1536, + 2048, + 3072, + 4096, + ] + else: + batch_sizes = [args.batch_size] + + ray.init() + num_gpus = int(ray.available_resources()["GPU"]) + workers = [BenchmarkWorker.remote(args.seed) for _ in range(num_gpus)] + + def _distribute(method: str, inputs: List[Any]) -> List[Any]: + outputs = [] + worker_idx = 0 + for input_args in inputs: + worker = workers[worker_idx] + worker_method = getattr(worker, method) + output = worker_method.remote(*input_args) + outputs.append(output) + worker_idx = (worker_idx + 1) % num_gpus + return ray.get(outputs) + + if args.tune: + search_space = get_configs_compute_bound() + if block_shape is not None: + block_n, block_k = block_shape[0], block_shape[1] + search_space = [ + config + for config in search_space + if block_k % config["BLOCK_SIZE_K"] == 0 + ] + print(f"Start tuning over {len(search_space)} configurations...") + + start = time.perf_counter() + configs = _distribute( + "tune", + [ + ( + batch_size, + E, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + search_space, + ) + for batch_size in batch_sizes + ], + ) + best_configs = { + M: sort_config(config) for M, config in zip(batch_sizes, configs) + } + save_configs( + best_configs, + E, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + ) + end = time.perf_counter() + print(f"Tuning took {end - start:.2f} seconds") + else: + outputs = _distribute( + "benchmark", + [ + ( + batch_size, + E, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + ) + for batch_size in batch_sizes + ], + ) + + for batch_size, (config, kernel_time) in zip(batch_sizes, outputs): + print(f"Batch size: {batch_size}, config: {config}") + print(f"Kernel time: {kernel_time:.2f} us") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" + ) + parser.add_argument("--tp-size", "--tp", type=int, default=2) + parser.add_argument( + "--dtype", + type=str, + choices=["auto", "fp8_w8a8", "int8_w8a16", "int8_w8a8"], + default="auto", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--batch-size", type=int, required=False) + parser.add_argument("--tune", action="store_true") + parser.add_argument("--disable-shared-experts-fusion", action="store_true") + args = parser.parse_args() + + main(args) diff --git a/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py b/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py new file mode 100644 index 000000000..78d81499e --- /dev/null +++ b/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py @@ -0,0 +1,576 @@ +import itertools +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange +from sgl_kernel import lightning_attention_decode as sgl_lightning_attention_decode + + +@triton.jit +def _decode_kernel( + Q, + K, + V, + KV, + Out, + S, + b: tl.constexpr, + h: tl.constexpr, + n: tl.constexpr, + d: tl.constexpr, + d_original: tl.constexpr, + e: tl.constexpr, + e_original: tl.constexpr, +): + off_bh = tl.program_id(0) + off_h = off_bh % h + + qk_offset = off_bh * n * d + v_offset = off_bh * n * e + o_offset = off_bh * n * e + kv_offset = off_bh * d * e + + s = tl.load(S + off_h) + ratio = tl.exp(-s) + + d_idx = tl.arange(0, d) + e_idx = tl.arange(0, e) + + # Create masks for original dimensions + d_mask = d_idx < d_original + e_mask = e_idx < e_original + + # Load with masking + q = tl.load(Q + qk_offset + d_idx, mask=d_mask, other=0.0) + k = tl.load(K + qk_offset + d_idx, mask=d_mask, other=0.0) + v = tl.load(V + v_offset + e_idx, mask=e_mask, other=0.0) + + # Load KV with 2D masking + kv = tl.load( + KV + kv_offset + d_idx[:, None] * e + e_idx[None, :], + mask=(d_mask[:, None] & e_mask[None, :]), + other=0.0, + ) + + # Compute outer product using element-wise operations + k_v_prod = k[:, None] * v[None, :] + kv = ratio * kv + k_v_prod + + # Store KV with 2D masking + tl.store( + KV + kv_offset + d_idx[:, None] * e + e_idx[None, :], + kv.to(KV.dtype.element_ty), + mask=(d_mask[:, None] & e_mask[None, :]), + ) + + # Compute matrix-vector multiplication using element-wise operations and reduction + o = tl.sum(q[:, None] * kv, axis=0) + + # Store output with masking + tl.store(Out + o_offset + e_idx, o.to(Out.dtype.element_ty), mask=e_mask) + + +def lightning_attn_decode(q, k, v, kv, s): + """Triton implementation of Lightning Attention decode operation""" + b, h, n, d = q.shape + e = v.shape[-1] + assert n == 1, "Sequence length must be 1 in decode mode" + + # Get padded dimensions (power of 2) + d_padded = next_power_of_2(d) + e_padded = next_power_of_2(e) + + # Create output tensor (padded) + o_padded = torch.empty(b, h, n, e_padded, dtype=v.dtype, device=v.device) + + # Create padded tensors without actually padding the data + q_padded = torch.empty(b, h, n, d_padded, dtype=q.dtype, device=q.device) + k_padded = torch.empty(b, h, n, d_padded, dtype=k.dtype, device=k.device) + v_padded = torch.empty(b, h, n, e_padded, dtype=v.dtype, device=v.device) + kv_padded = torch.empty( + b, h, d_padded, e_padded, dtype=torch.float32, device=kv.device + ) + + # Copy data to padded tensors + q_padded[..., :d] = q + k_padded[..., :d] = k + v_padded[..., :e] = v + kv_padded[..., :d, :e] = kv + + # Launch kernel + grid = (b * h, 1) + _decode_kernel[grid]( + q_padded, + k_padded, + v_padded, + kv_padded, + o_padded, + s, + b=b, + h=h, + n=n, + d=d_padded, + d_original=d, + e=e_padded, + e_original=e, + ) + + # Get unpadded outputs + o = o_padded[..., :e] + kv_out = kv_padded[..., :d, :e] + + return o, kv_out + + +def next_power_of_2(n): + return 2 ** (int(math.ceil(math.log(n, 2)))) + + +class MiniMaxText01LightningAttention(nn.Module): + def __init__(self, config=None, layer_idx: Optional[int] = None, **kwargs): + super().__init__() + if config is None: + config = type("Config", (), kwargs) + + bias = False + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) + + self.out_proj = nn.Linear( + self.head_dim * self.num_heads, self.hidden_size, bias=bias + ) + self.act = get_activation_fn(config.hidden_act) + self.norm = MiniMaxText01RMSNorm(self.head_dim * self.num_heads) + + self.qkv_proj = nn.Linear( + self.hidden_size, 3 * self.head_dim * self.num_heads, bias=bias + ) + self.output_gate = nn.Linear( + self.hidden_size, self.head_dim * self.num_heads, bias=bias + ) + + # for inference only + self.offset = 0 + self.layer_idx = layer_idx + + def forward( + self, + hidden_states, + attn_mask: Optional[torch.Tensor] = None, # (b, h, n, m) + output_attentions: bool = False, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + use_cache: bool = False, + slope_rate: Optional[torch.Tensor] = None, + **kwargs, + ): + if (not self.training) and (not do_eval): + return self.inference( + hidden_states, + attn_mask, + output_attentions, + past_key_value, + use_cache, + slope_rate, + ) + + def inference( + self, + x, + attn_mask: Optional[torch.Tensor] = None, # (b, n) + output_attentions: bool = False, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + use_cache: bool = False, + slope_rate: Optional[torch.Tensor] = None, # (h, 1, 1) + ): + # x: b n d + b, n, d = x.shape + # linear map + qkv = self.act(self.qkv_proj(x)) + new_shape = qkv.size()[:-1] + (self.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [self.head_dim] * 3, dim=3) + q = q.transpose(1, 2) # [b, n, h, d] -> [b, h, n, d] + k = k.transpose(1, 2) # [b, n, h, d] -> [b, h, n, d] + v = v.transpose(1, 2) # [b, n, h, d] -> [b, h, n, e] + + self.offset += 1 + ratio = torch.exp(-slope_rate) # [h, 1, 1] + + # decode mode + kv = past_key_value # [b, h, d, e] + output = [] + for i in range(n): + # kv: [b, h, d, e] + # ratio: [h, 1, 1] + # k: [b, h, n, d] + # v: [b, h, n, e] + # k[:, :, i : i + 1]: [b, h, 1, d] + # v[:, :, i : i + 1]: [b, h, 1, e] + # ratio * kv: [b, h, d, e] + # torch.einsum( + # "... n d, ... n e -> ... d e", + # k[:, :, i : i + 1], + # v[:, :, i : i + 1], + # ) + # [b, h, d, e] + [b, h, d, e] -> [b, h, d, e] + kv = ratio * kv + torch.einsum( + "... n d, ... n e -> ... d e", + k[:, :, i : i + 1], + v[:, :, i : i + 1], + ) + # q[:, :, i : i + 1]: [b, h, 1, d] + # kv.to(q.dtype): [b, h, d, e] + # torch.einsum( + # "... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv.to(q.dtype) + # ) + # [b, h, 1, d] * [b, h, d, e] -> [b, h, 1, e] + qkv = torch.einsum( + "... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv.to(q.dtype) + ) + output.append(qkv) + output = torch.cat(output, dim=-2) + + # reshape + output = rearrange(output, "b h n d -> b n (h d)") + # normalize + output = self.norm(output) + # gate + output = F.sigmoid(self.output_gate(x)) * output + # outproj + output = self.out_proj(output) + + attn_weights = None + + return output, attn_weights, kv + + +def get_activation_fn(activation): + if activation == "gelu": + return F.gelu + elif activation == "relu": + return F.relu + elif activation == "elu": + return F.elu + elif activation == "sigmoid": + return F.sigmoid + elif activation == "exp": + + def f(x): + with torch.no_grad(): + x_max = torch.max(x, dim=-1, keepdims=True).values + y = torch.exp(x - x_max) + return y + + return f + elif activation == "leak": + return F.leaky_relu + elif activation == "1+elu": + + def f(x): + return 1 + F.elu(x) + + return f + elif activation == "2+elu": + + def f(x): + return 2 + F.elu(x) + + return f + elif activation == "silu" or activation == "swish": + return F.silu + elif activation == "sine": + return torch.sin + else: + return lambda x: x + + +class MiniMaxText01RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + MiniMaxText01RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +def test_lightning_attention_implementations(model_params): + torch.manual_seed(42) + + batch_size = 64 + seq_len = 1 + dtype = torch.bfloat16 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + hidden_states = torch.randn( + batch_size, seq_len, model_params["hidden_size"], dtype=dtype, device=device + ) + + attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device) + + slope_rate = _build_slope_tensor(model_params["num_attention_heads"]).to(device) + + model_attn = MiniMaxText01LightningAttention(**model_params).to(dtype).to(device) + model_attn.eval() + + d = model_params["head_dim"] + past_kv = torch.randn( + batch_size, + model_params["num_attention_heads"], + d, + d, + device=device, + ) + with torch.no_grad(): + model_output, _, new_kv = model_attn.inference( + hidden_states, + attn_mask=attention_mask, + slope_rate=slope_rate, + past_key_value=past_kv, + ) + + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + past_kv = past_kv.contiguous() + slope_rate = slope_rate.contiguous() + + # Test Triton implementation + triton_output, triton_new_kv = lightning_attn_decode(q, k, v, past_kv, slope_rate) + triton_output = triton_output.transpose(1, 2).contiguous() + triton_output = triton_output.view(batch_size, seq_len, -1) + triton_output = model_attn.norm(triton_output) + triton_output = torch.sigmoid(model_attn.output_gate(hidden_states)) * triton_output + triton_output = model_attn.out_proj(triton_output) + + # Test SGL implementation + sgl_output = torch.empty_like(v) + sgl_new_kv = torch.empty_like(past_kv) + sgl_lightning_attention_decode(q, k, v, past_kv, slope_rate, sgl_output, sgl_new_kv) + + sgl_output = sgl_output.transpose(1, 2).contiguous() + sgl_output = sgl_output.view(batch_size, seq_len, -1) + sgl_output = model_attn.norm(sgl_output) + sgl_output = torch.sigmoid(model_attn.output_gate(hidden_states)) * sgl_output + sgl_output = model_attn.out_proj(sgl_output) + + # Verify Triton implementation results + torch.testing.assert_close( + model_output, + triton_output, + rtol=1e-3, + atol=1e-2, + msg="Triton lightning attention implementation produces different output results", + ) + torch.testing.assert_close( + new_kv, + triton_new_kv, + rtol=1e-3, + atol=1e-2, + msg="Triton lightning attention implementation produces different kv results", + ) + + # Verify SGL implementation results + torch.testing.assert_close( + model_output, + sgl_output, + rtol=1e-3, + atol=1e-2, + msg="SGL lightning attention implementation produces different output results", + ) + torch.testing.assert_close( + new_kv, + sgl_new_kv, + rtol=1e-3, + atol=1e-2, + msg="SGL lightning attention implementation produces different kv results", + ) + + print("✅ All implementations match") + + +def _build_slope_tensor(n_attention_heads: int): + def get_slopes(n): + def get_slopes_power_of_2(n): + start = 2 ** (-(2 ** -(math.log2(n) - 3))) + ratio = start + return [start * ratio**i for i in range(n)] + + if math.log2(n).is_integer(): + return get_slopes_power_of_2(n) + else: + closest_power_of_2 = 2 ** math.floor(math.log2(n)) + return ( + get_slopes_power_of_2(closest_power_of_2) + + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2] + ) + + slopes = torch.tensor(get_slopes(n_attention_heads)).reshape( + n_attention_heads, 1, 1 + ) + return slopes + + +def get_benchmark(): + batch_size_range = [i for i in range(1, 33)] # max 32 + seq_length_range = [1] # decode mode sequence length is fixed to 1 + configs = list(itertools.product(batch_size_range, seq_length_range)) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "seq_len"], + x_vals=[list(_) for _ in configs], + line_arg="provider", + line_vals=["Original", "Triton", "SGL"], + line_names=[ + "Original PyTorch Implementation", + "Triton Implementation", + "SGL Implementation", + ], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="us", + plot_name="lightning-attention-decode-performance", + args={}, + ) + ) + def benchmark(batch_size, seq_len, provider): + dtype = torch.bfloat16 + device = torch.device("cuda") + + params = { + "hidden_size": 6144, + "num_attention_heads": 64, + "head_dim": 96, + "hidden_act": "gelu", + } + + hidden_states = torch.randn( + batch_size, seq_len, params["hidden_size"], dtype=dtype, device=device + ) + + attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device) + + slope_rate = _build_slope_tensor(params["num_attention_heads"]).to(device) + model_attn = MiniMaxText01LightningAttention(**params).to(dtype).to(device) + model_attn.eval() + + d = params["head_dim"] + past_kv = torch.randn( + batch_size, + params["num_attention_heads"], + d, + d, + device=device, + ) + + quantiles = [0.5, 0.2, 0.8] + if provider == "Original": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: model_attn.inference( + hidden_states, + attn_mask=attention_mask, + slope_rate=slope_rate, + past_key_value=past_kv, + ), + quantiles=quantiles, + ) + elif provider == "Triton": + + def run_triton(): + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + output, new_kv = lightning_attn_decode(q, k, v, past_kv, slope_rate) + output = output.transpose(1, 2).contiguous() + output = output.view(batch_size, seq_len, -1) + output = model_attn.norm(output) + output = torch.sigmoid(model_attn.output_gate(hidden_states)) * output + return model_attn.out_proj(output) + + ms, min_ms, max_ms = triton.testing.do_bench( + run_triton, + quantiles=quantiles, + ) + else: # SGL + + def run_sgl(): + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + v = v.transpose(1, 2).contiguous() + + output = torch.empty_like(v) + new_kv = torch.empty_like(past_kv) + sgl_lightning_attention_decode( + q, k, v, past_kv, slope_rate, output, new_kv + ) + + output = output.transpose(1, 2).contiguous() + output = output.view(batch_size, seq_len, -1) + output = model_attn.norm(output) + output = torch.sigmoid(model_attn.output_gate(hidden_states)) * output + return model_attn.out_proj(output) + + ms, min_ms, max_ms = triton.testing.do_bench( + run_sgl, + quantiles=quantiles, + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/lightning_attention_decode/", + help="Path to save lightning attention decode benchmark results", + ) + args = parser.parse_args() + + params = { + "hidden_size": 6144, + "num_attention_heads": 64, + "head_dim": 96, + "hidden_act": "silu", + } + # Run correctness test first + # Adapted from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/config.json + test_lightning_attention_implementations(params) + + # Run performance benchmark + benchmark = get_benchmark() + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py b/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py new file mode 100644 index 000000000..3bf9054bd --- /dev/null +++ b/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py @@ -0,0 +1,603 @@ +import itertools +import math +import os +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange + + +# Adapted from https://github.com/OpenNLPLab/lightning-attention/blob/main/lightning_attn/ops/triton/lightning_attn2.py +@triton.jit +def _fwd_kernel( + Q, + K, + V, + Out, + S, # log lambda + b: tl.constexpr, + h: tl.constexpr, + n: tl.constexpr, + d: tl.constexpr, + e: tl.constexpr, + BLOCK: tl.constexpr, + NUM_BLOCK: tl.constexpr, + BLOCK_MODEL: tl.constexpr, +): + ##### get offset + off_bh = tl.program_id(0) + off_h = off_bh % h + off_e = tl.program_id(1) + qk_offset = off_bh * n * d + v_offset = off_bh * n * e + o_offset = off_bh * n * e + # channel offset + e_offset = off_e * BLOCK_MODEL + + ##### get block ptr + Q_block_ptr = Q + qk_offset + tl.arange(0, d)[None, :] + K_trans_block_ptr = K + qk_offset + tl.arange(0, d)[:, None] + V_block_ptr = V + v_offset + e_offset + tl.arange(0, BLOCK_MODEL)[None, :] + O_block_ptr = Out + o_offset + e_offset + tl.arange(0, BLOCK_MODEL)[None, :] + S_block_ptr = S + off_h + + ##### init diag decay(Lambda); q, k decay; kv + s = tl.load(S_block_ptr) + # q, k decay + off_block = tl.arange( + 0, BLOCK + ) # Not bug, this is a bit different from algorithm 1, but is mathematically equivalent + q_decay = tl.exp(-s.to(tl.float32) * off_block[:, None]) + k_trans_decay = tl.exp(-s.to(tl.float32) * (BLOCK - off_block[None, :])) + block_decay = tl.exp(-s.to(tl.float32) * BLOCK) + # diag decay + index = off_block[:, None] - off_block[None, :] + s_index = s * index + s_index = tl.where(index >= 0, -s_index, float("-inf")) + diag_decay = tl.exp(s_index) + kv = tl.zeros([d, BLOCK_MODEL], dtype=tl.float32) + + ##### compute + for i in range(NUM_BLOCK): + # load + q = tl.load( + Q_block_ptr + off_block[:, None] * d, mask=off_block[:, None] < n, other=0.0 + ).to(tl.float32) + k_trans = tl.load( + K_trans_block_ptr + off_block[None, :] * d, + mask=off_block[None, :] < n, + other=0.0, + ).to(tl.float32) + v = tl.load( + V_block_ptr + off_block[:, None] * e, mask=off_block[:, None] < n, other=0.0 + ).to(tl.float32) + + # compute + qk = tl.dot(q, k_trans) * diag_decay + o_intra = tl.dot(qk, v) + o_inter = tl.dot(q, kv) * q_decay + o = o_intra + o_inter + + # save and update + tl.store( + O_block_ptr + off_block[:, None] * e, + o.to(O_block_ptr.dtype.element_ty), + mask=off_block[:, None] < n, + ) + kv = block_decay * kv + tl.dot(k_trans * k_trans_decay, v) + off_block += BLOCK + + +def lightning_attn2(q, k, v, s): + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + s = s.contiguous() + + b, h, n, d = q.shape + e = v.shape[-1] + + # Pad d to next power of 2 + d_padded = next_power_of_2(d) + if d_padded != d: + q_padded = F.pad(q, (0, d_padded - d)) + k_padded = F.pad(k, (0, d_padded - d)) + else: + q_padded = q + k_padded = k + + # Pad e to next power of 2 + e_padded = next_power_of_2(e) + if e_padded != e: + v_padded = F.pad(v, (0, e_padded - e)) + else: + v_padded = v + + o_padded = torch.empty((b, h, n, e_padded), dtype=q.dtype, device=q.device) + + BLOCK = 64 + NUM_BLOCK = triton.cdiv(q.shape[2], BLOCK) + # parallel over channel + BLOCK_MODEL = min(triton.next_power_of_2(e_padded), 32) + grid = (b * h, triton.cdiv(e_padded, BLOCK_MODEL)) + + _fwd_kernel[grid]( + q_padded, + k_padded, + v_padded, + o_padded, + s, + b, + h, + n, + d_padded, + e_padded, + BLOCK=BLOCK, + NUM_BLOCK=NUM_BLOCK, + BLOCK_MODEL=BLOCK_MODEL, + ) + + # Remove padding from output + if e_padded != e: + o = o_padded[..., :e] + else: + o = o_padded + + return o + + +def is_support(dim): + return 16 % dim + + +def next_power_of_2(n): + return 2 ** (int(math.ceil(math.log(n, 2)))) + + +def lightning_attn_func(q, k, v, s): + b, h, n, d = q.shape + e = v.shape[-1] + assert is_support(d) and is_support(e) + + # pad v's feature dim to power of 2 + e_pad = next_power_of_2(e) + need_pad = e_pad != e + if need_pad: + v = F.pad(v, (0, e_pad - e)) + + if d > 128: + # split over head + if 64 % d: + m = 64 + elif 32 % d: + m = 32 + elif 16 % d: + m = 16 + arr = [m * i for i in range(d // m + 1)] + if arr[-1] != d: + arr.append(d) + n = len(arr) + o = 0 + for i in range(n - 1): + start = arr[i] + end = arr[i + 1] + q1 = q[..., start:end] + k1 = k[..., start:end] + o += lightning_attn2(q1, k1, v, s) + else: + o = lightning_attn2(q, k, v, s) + + if need_pad: + o = o[:, :, :, :e] + + return o + + +debug = eval(os.environ.get("debug", default="False")) + +BLOCK = 256 + + +# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->MiniMaxText01 +class MiniMaxText01RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + MiniMaxText01RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +# Copied from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/modeling_minimax_text_01.py +def get_activation_fn(activation): + if debug: + logger.info(f"activation: {activation}") + if activation == "gelu": + return F.gelu + elif activation == "relu": + return F.relu + elif activation == "elu": + return F.elu + elif activation == "sigmoid": + return F.sigmoid + elif activation == "exp": + + def f(x): + with torch.no_grad(): + x_max = torch.max(x, dim=-1, keepdims=True).values + y = torch.exp(x - x_max) + + return y + + return f + elif activation == "leak": + return F.leaky_relu + elif activation == "1+elu": + + def f(x): + return 1 + F.elu(x) + + return f + elif activation == "2+elu": + + def f(x): + return 2 + F.elu(x) + + return f + elif activation == "silu" or activation == "swish": + return F.silu + elif activation == "sine": + return torch.sin + else: + logger.info(f"activation: does not support {activation}, use Identity!!!") + return lambda x: x + + +# Copied from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/modeling_minimax_text_01.py +class MiniMaxText01LightningAttention(nn.Module): + def __init__(self, config=None, layer_idx: Optional[int] = None, **kwargs): + super().__init__() + if config is None: + config = type("Config", (), kwargs) + + bias = False + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) + + self.out_proj = nn.Linear( + self.head_dim * self.num_heads, self.hidden_size, bias=bias + ) + self.act = get_activation_fn(config.hidden_act) + self.norm = MiniMaxText01RMSNorm(self.head_dim * self.num_heads) + + self.qkv_proj = nn.Linear( + self.hidden_size, 3 * self.head_dim * self.num_heads, bias=bias + ) + self.output_gate = nn.Linear( + self.hidden_size, self.head_dim * self.num_heads, bias=bias + ) + + # for inference only + self.offset = 0 + self.layer_idx = layer_idx + + def forward( + self, + hidden_states, + attn_mask: Optional[torch.Tensor] = None, # (b, h, n, m) + output_attentions: bool = False, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + use_cache: bool = False, + slope_rate: Optional[torch.Tensor] = None, + **kwargs, + ): + if (not self.training) and (not do_eval): + return self.inference( + hidden_states, + attn_mask, + output_attentions, + past_key_value, + use_cache, + slope_rate, + ) + + def inference( + self, + x, + attn_mask: Optional[torch.Tensor] = None, # (b, n) + output_attentions: bool = False, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + use_cache: bool = False, + slope_rate: Optional[torch.Tensor] = None, # (h, 1, 1) + ): + # x: b n d + b, n, d = x.shape + # linear map + qkv = self.act(self.qkv_proj(x)) + new_shape = qkv.size()[:-1] + (self.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [self.head_dim] * 3, dim=3) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + if past_key_value is None: + self.offset = q.shape[-2] + else: + self.offset += 1 + + # for align with metaseq + ratio = torch.exp(-slope_rate) + + # only use for the first time + if past_key_value is None: + slope_rate = slope_rate.to(torch.float32) + if attn_mask is not None: + v = v.masked_fill( + (1 - attn_mask).unsqueeze(1).unsqueeze(-1).to(torch.bool), 0 + ) + NUM_BLOCK = (n + BLOCK - 1) // BLOCK + b, h, n, d = q.shape + e = v.shape[-1] + # other + array = torch.arange(BLOCK).to(q) + 1 + q_decay = torch.exp(-slope_rate * array.reshape(-1, 1)) + k_decay = torch.exp(-slope_rate * (BLOCK - array.reshape(-1, 1))) + index = array[:, None] - array[None, :] + s_index = ( + slope_rate + * index[ + None, + None, + ] + ) + s_index = torch.where(index >= 0, -s_index, float("-inf")) + diag_decay = torch.exp(s_index) + + kv = torch.zeros(b, h, d, e).to(torch.float32).to(q.device) + output = torch.empty((b, h, n, e), dtype=q.dtype, device=q.device) + for i in range(NUM_BLOCK): + si = i * BLOCK + ei = min(si + BLOCK, n) + m = ei - si + qi = q[:, :, si:ei].contiguous() + ki = k[:, :, si:ei].contiguous() + vi = v[:, :, si:ei].contiguous() + qkv_none_diag = torch.matmul(qi * q_decay[:, :m], kv).to(torch.float32) + + # diag + qk = ( + torch.matmul(qi, ki.transpose(-1, -2)).to(torch.float32) + * diag_decay[:, :, :m, :m] + ) + qkv_diag = torch.matmul(qk, vi.to(torch.float32)) + block_decay = torch.exp(-slope_rate * m) + output[:, :, si:ei] = qkv_none_diag + qkv_diag + kv = block_decay * kv + torch.matmul( + (ki * k_decay[:, -m:]).transpose(-1, -2).to(vi.dtype), vi + ) + + else: + kv = past_key_value + output = [] + for i in range(n): + kv = ratio * kv + torch.einsum( + "... n d, ... n e -> ... d e", + k[:, :, i : i + 1], + v[:, :, i : i + 1], + ) + qkv = torch.einsum( + "... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv.to(q.dtype) + ) + output.append(qkv) + output = torch.cat(output, dim=-2) + # reshape + output = rearrange(output, "b h n d -> b n (h d)") + # normalize + output = self.norm(output) + # gate + output = F.sigmoid(self.output_gate(x)) * output + # outproj + output = self.out_proj(output) + + attn_weights = None + + return output, attn_weights, kv + + +def _build_slope_tensor(n_attention_heads: int): + def get_slopes(n): + def get_slopes_power_of_2(n): + start = 2 ** (-(2 ** -(math.log2(n) - 3))) + ratio = start + return [start * ratio**i for i in range(n)] + + if math.log2(n).is_integer(): + return get_slopes_power_of_2( + n + ) # In the paper, we only train models that have 2^a heads for some a. This function has + else: # some good properties that only occur when the input is a power of 2. To maintain that even + closest_power_of_2 = 2 ** math.floor( + math.log2(n) + ) # when the number of heads is not a power of 2, we use this workaround. + return ( + get_slopes_power_of_2(closest_power_of_2) + + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2] + ) + + # h, 1, 1 + slopes = torch.tensor(get_slopes(n_attention_heads)).reshape( + n_attention_heads, 1, 1 + ) + + return slopes + + +def test_lightning_attention_implementations(model_params): + torch.manual_seed(42) + + batch_size = 2 + seq_len = 1024 + dtype = torch.bfloat16 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + hidden_states = torch.randn( + batch_size, seq_len, model_params["hidden_size"], dtype=dtype, device=device + ) + + attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device) + + slope_rate = _build_slope_tensor(model_params["num_attention_heads"]).to(device) + + model_attn = MiniMaxText01LightningAttention(**model_params).to(dtype).to(device) + model_attn.eval() + + with torch.no_grad(): + model_output, _, _ = model_attn.inference( + hidden_states, attn_mask=attention_mask, slope_rate=slope_rate + ) + + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + lib_output = lightning_attn_func(q, k, v, slope_rate) + lib_output = lib_output.transpose(1, 2).contiguous() + lib_output = lib_output.view(batch_size, seq_len, -1) + lib_output = model_attn.norm(lib_output) + lib_output = torch.sigmoid(model_attn.output_gate(hidden_states)) * lib_output + lib_output = model_attn.out_proj(lib_output) + + torch.testing.assert_close( + model_output, + lib_output, + rtol=1e-3, + atol=1e-2, + msg="Lightning attention implementations produce different results", + ) + + print("✅ Two implementations match") + + +def get_benchmark(): + batch_size_range = [2**i for i in range(0, 7)] # max 64 + seq_length_range = [256, 512, 1024, 2048, 4096] # max 4096 + configs = list(itertools.product(batch_size_range, seq_length_range)) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "seq_len"], + x_vals=[list(_) for _ in configs], + line_arg="provider", + line_vals=["MiniMax-Text-01", "OpenNLPLab"], + line_names=[ + "MiniMax-Text-01 Model Implementation", + "OpenNLPLab Library Implementation", + ], + styles=[("blue", "-"), ("green", "-")], + ylabel="us", + plot_name="lightning-attention-prefill-performance", + args={}, + ) + ) + def benchmark(batch_size, seq_len, provider): + dtype = torch.bfloat16 + device = torch.device("cuda") + + params = { + "hidden_size": 6144, + "num_attention_heads": 64, + "head_dim": 96, + "hidden_act": "gelu", + } + + hidden_states = torch.randn( + batch_size, seq_len, params["hidden_size"], dtype=dtype, device=device + ) + + attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device) + + slope_rate = _build_slope_tensor(params["num_attention_heads"]).to(device) + model_attn = MiniMaxText01LightningAttention(**params).to(dtype).to(device) + model_attn.eval() + + quantiles = [0.5, 0.2, 0.8] + if provider == "MiniMax-Text-01": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: model_attn.inference( + hidden_states, attn_mask=attention_mask, slope_rate=slope_rate + ), + quantiles=quantiles, + ) + else: + + def run_lib(): + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + lib_output = lightning_attn_func(q, k, v, slope_rate) + lib_output = lib_output.transpose(1, 2).contiguous() + lib_output = lib_output.view(batch_size, seq_len, -1) + lib_output = model_attn.norm(lib_output) + lib_output = ( + torch.sigmoid(model_attn.output_gate(hidden_states)) * lib_output + ) + return model_attn.out_proj(lib_output) + + ms, min_ms, max_ms = triton.testing.do_bench( + run_lib, + quantiles=quantiles, + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/lightning_attention_prefill/", + help="Path to save lightning attention prefill benchmark results", + ) + args = parser.parse_args() + + # Run correctness test first + # Adapted from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/config.json + params = { + "hidden_size": 6144, + "num_attention_heads": 64, + "head_dim": 96, + "hidden_act": "silu", + } + test_lightning_attention_implementations(params) + + # Run performance benchmark + benchmark = get_benchmark() + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/benchmark/kernels/quantization/bench_fp4_quant.py b/benchmark/kernels/quantization/bench_fp4_quant.py new file mode 100644 index 000000000..318e820ad --- /dev/null +++ b/benchmark/kernels/quantization/bench_fp4_quant.py @@ -0,0 +1,133 @@ +import argparse +import itertools + +import torch +import triton +from sgl_kernel import scaled_fp4_grouped_quant, silu_and_mul_scaled_fp4_grouped_quant +from sgl_kernel.elementwise import silu_and_mul + +from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_masked_post_quant_fwd +from sglang.srt.layers.quantization import deep_gemm_wrapper + + +def _test_accuracy_once(E, M, K, input_dtype, device): + x = torch.randn(E, M, K, device=device, dtype=input_dtype) + glb_scales = torch.ones((E,), dtype=torch.float32, device=device) + masks = torch.full((E,), M, dtype=torch.int32, device=device) + out, blk_scales = silu_and_mul_scaled_fp4_grouped_quant(x, glb_scales, masks) + out1, blk_scales1 = scaled_fp4_grouped_quant( + silu_and_mul(x), + glb_scales, + masks, + ) + + torch.testing.assert_close(out, out1) + torch.testing.assert_close(blk_scales, blk_scales1) + print(f"E: {E}, M: {M}, K: {K}, type: {input_dtype} OK") + + +NUM_RANKS = 48 +M_PER_RANKs = [128, 256, 512, 1024] +Ms = [M_PER_RANK * NUM_RANKS for M_PER_RANK in M_PER_RANKs] +Ks = [2048, 4096, 7168] + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["M", "K"], + x_vals=list(itertools.product(Ms, Ks)), + x_log=False, + line_arg="provider", + line_vals=["triton_fp8", "cuda_unfused_fp4", "cuda_fused_fp4"], + line_names=["triton_fp8", "cuda_unfused_fp4", "cuda_fused_fp4"], + styles=[("blue", "-"), ("orange", "-"), ("green", "-")], + ylabel="ms", + plot_name="fp4 quant", + args={}, + ) +) +def benchmark(M, K, provider): + E = 6 + device = "cuda" + x = torch.randn(E, M, K, device=device, dtype=torch.bfloat16) + glb_scales = torch.ones((E,), dtype=torch.float32, device=device) + masks = torch.randint(1, 4096, (E,), dtype=torch.int32, device=device) + fp8_out = torch.empty( + ( + x.shape[0], + x.shape[1], + x.shape[2] // 2, + ), + device=x.device, + dtype=torch.float8_e4m3fn, + ) + scale_block_size = 128 + fp8_scales = torch.empty( + ( + x.shape[0], + x.shape[1], + x.shape[2] // 2 // scale_block_size, + ), + device=x.device, + dtype=torch.float32, + ) + + quantiles = [0.5, 0.2, 0.8] + if provider == "triton_fp8": + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + lambda: silu_and_mul_masked_post_quant_fwd( + x, + fp8_out, + fp8_scales, + scale_block_size, + masks, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ), + quantiles=quantiles, + ) + if provider == "cuda_unfused_fp4": + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + lambda: scaled_fp4_grouped_quant( + silu_and_mul(x), + glb_scales, + masks, + ), + quantiles=quantiles, + ) + if provider == "cuda_fused_fp4": + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + lambda: silu_and_mul_scaled_fp4_grouped_quant( + x, + glb_scales, + masks, + ), + quantiles=quantiles, + ) + + return ms, min_ms, max_ms + + +def test_accuracy(): + E = 6 + N_RANKS = 48 + Ms = [128, 256, 512, 1024] + Ks = [2048, 4096, 7168] + input_dtype = torch.bfloat16 + for M in Ms: + for K in Ks: + _test_accuracy_once(E, N_RANKS * M, K, input_dtype, "cuda") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./bench_fp4_quant_res", + help="Path to save fp4 quant benchmark results", + ) + args = parser.parse_args() + + test_accuracy() + + benchmark.run(print_data=True, show_plots=True, save_path=args.save_path) diff --git a/benchmark/kernels/quantization/bench_int8_quant.py b/benchmark/kernels/quantization/bench_int8_quant.py new file mode 100644 index 000000000..94b795690 --- /dev/null +++ b/benchmark/kernels/quantization/bench_int8_quant.py @@ -0,0 +1,94 @@ +import argparse + +import torch +import triton +from vllm._custom_ops import scaled_int8_quant as vllm_scaled_int8_quant + +from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8 + + +@torch.compile(backend="inductor") +def torch_int8_quant(x): + int8_max = torch.iinfo(torch.int8).max + + abs_max = x.abs().max(dim=-1, keepdim=True).values + scales = abs_max.to(torch.float32) / float(int8_max) + + q_x = (x / scales).round().to(torch.int8) + + return q_x, scales + + +def _test_accuracy_once(M, K, input_dtype, device): + x = torch.randn(M, K, dtype=input_dtype, device=device) * 5000 + out, scales, _ = vllm_scaled_int8_quant(x, symmetric=True) + out1, scales1 = per_token_quant_int8(x) + out2, scales2 = torch_int8_quant(x) + torch.testing.assert_close(out, out2, atol=1, rtol=0) + torch.testing.assert_close(out, out1, atol=1, rtol=0) + torch.testing.assert_close(scales, scales2) + torch.testing.assert_close(scales1, scales2) + print(f"M: {M}, K: {K}, type: {input_dtype} OK") + + +def test_accuracy(): + Ms = [1, 13, 128, 1024, 2048, 4096] + Ks = [512, 1024, 2048, 8192] + input_dtypes = [torch.float16, torch.bfloat16] + for M in Ms: + for K in Ks: + for input_dtype in input_dtypes: + _test_accuracy_once(M, K, input_dtype, "cuda") + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048], + x_log=False, + line_arg="provider", + line_vals=["vllm op", "triton", "torch.compile"], + line_names=["vllm op", "triton", "torch.compile"], + styles=[("blue", "-"), ("orange", "-"), ("red", "-")], + ylabel="ms", + plot_name="int8 per token quant", + args={}, + ) +) +def benchmark(batch_size, provider): + M, K = batch_size, 16384 + x = torch.randn(M, K, dtype=torch.float16, device="cuda") * 1000 + + quantiles = [0.5, 0.2, 0.8] + if provider == "vllm op": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: vllm_scaled_int8_quant(x, symmetric=True), + quantiles=quantiles, + ) + if provider == "triton": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: per_token_quant_int8(x), + quantiles=quantiles, + ) + if provider == "torch.compile": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: torch_int8_quant(x), + quantiles=quantiles, + ) + + return ms, min_ms, max_ms + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./bench_int8_quant_res", + help="Path to save int8 quant benchmark results", + ) + args = parser.parse_args() + + test_accuracy() + + benchmark.run(print_data=True, show_plots=True, save_path=args.save_path) diff --git a/benchmark/kernels/quantization/tuning_block_wise_kernel.py b/benchmark/kernels/quantization/tuning_block_wise_kernel.py new file mode 100644 index 000000000..1b51e54b7 --- /dev/null +++ b/benchmark/kernels/quantization/tuning_block_wise_kernel.py @@ -0,0 +1,474 @@ +# Copyright 2025 SGLang 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 json +import multiprocessing as mp +import os +import time +from datetime import datetime +from typing import Any, Dict, List + +import torch +import triton +from tqdm import tqdm + +mp.set_start_method("spawn", force=True) + +from sglang.srt.layers.quantization.fp8_kernel import ( + _w8a8_block_fp8_matmul, + _w8a8_block_fp8_matmul_unrolledx4, +) +from sglang.srt.layers.quantization.int8_kernel import _w8a8_block_int8_matmul +from sglang.srt.utils import get_device_core_count, get_device_name, is_hip + +_is_hip = is_hip() + +DTYPE_MAP = { + "float32": torch.float32, + "float16": torch.float16, + "half": torch.half, + "bfloat16": torch.bfloat16, +} + + +def w8a8_block_matmul( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + config: Dict[str, Any], + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + """This function performs matrix multiplication with block-wise quantization. + + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + + Returns: + torch.Tensor: The result of matmul. + """ + assert len(block_size) == 2 + block_n, block_k = block_size[0], block_size[1] + + assert A.shape[-1] == B.shape[-1] + assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() + assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] + M = A.numel() // A.shape[-1] + + assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 + N, K = B.shape + assert triton.cdiv(N, block_n) == Bs.shape[0] + assert triton.cdiv(K, block_k) == Bs.shape[1] + + C_shape = A.shape[:-1] + (N,) + C = A.new_empty(C_shape, dtype=output_dtype) + + def grid(META): + return ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + + # Use manually unrolledx4 kernel on AMD GPU when the grid size is small. + # Empirical testing shows the sweet spot lies when it's less than the # of + # compute units available on the device. + num_workgroups = triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv( + N, config["BLOCK_SIZE_N"] + ) + + if A.dtype == torch.float8_e4m3fnuz or A.dtype == torch.float8_e4m3fn: + kernel = ( + _w8a8_block_fp8_matmul_unrolledx4 + if (_is_hip == True and num_workgroups <= get_device_core_count()) + else _w8a8_block_fp8_matmul + ) + else: + kernel = _w8a8_block_int8_matmul + + kernel[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + **config, + ) + + return C + + +def get_rocm_configs_compute_bound(): + configs = [] + waves_per_eu_range = 0 + for num_stages in [2]: + for block_m in [32, 64, 128, 256]: + for block_k in [32, 64, 128, 256]: + for block_n in [16, 32, 64, 128, 256]: + for num_warps in [4, 8]: + for group_size in [1, 4, 8, 16, 32]: + configs.append( + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_size, + "num_warps": num_warps, + "num_stages": num_stages, + "waves_per_eu": waves_per_eu_range, + } + ) + return configs + + +def get_configs_compute_bound(): + configs = [] + if _is_hip: + configs = get_rocm_configs_compute_bound() + else: + for num_stages in [2, 3, 4, 5]: + for block_m in [16, 32, 64, 128, 256]: + for block_k in [64, 128]: + for block_n in [32, 64, 128, 256]: + for num_warps in [4, 8]: + for group_size in [1, 16, 32, 64]: + configs.append( + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_size, + "num_warps": num_warps, + "num_stages": num_stages, + } + ) + return configs + + +def get_weight_shapes(tp_size): + # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3. Modify them, if you tune for another different model. + # cannot TP + total = [ + (512 + 64, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (7168, 16384), + (7168, 18432), + ] + # N can TP + n_tp = [ + (18432 * 2, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (24576, 1536), + (4096, 7168), + ] + # K can TP + k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] + + weight_shapes = [] + for t in total: + weight_shapes.append(t) + for n_t in n_tp: + new_t = (n_t[0] // tp_size, n_t[1]) + weight_shapes.append(new_t) + for k_t in k_tp: + new_t = (k_t[0], k_t[1] // tp_size) + weight_shapes.append(new_t) + return weight_shapes + + +def benchmark_config( + A, B, As, Bs, block_size, config, out_dtype=torch.float16, num_iters=10 +): + def run(): + w8a8_block_matmul(A, B, As, Bs, block_size, config, out_dtype) + + torch.cuda.synchronize() + # JIT complication & warmup + for _ in range(5): + run() + torch.cuda.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + latencies: List[float] = [] + for i in range(num_iters): + torch.cuda.synchronize() + start_event.record() + run() + end_event.record() + end_event.synchronize() + latencies.append(start_event.elapsed_time(end_event)) + avg = sum(latencies) / (num_iters * 10) * 1000 # us + return avg + + +def tune(M, N, K, block_size, out_dtype, search_space, input_type): + factor_for_scale = 1e-2 + + if input_type == "fp8": + fp8_info = torch.finfo( + torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn + ) + fp8_max, fp8_min = fp8_info.max, fp8_info.min + + A_fp32 = ( + (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max + ) + A = A_fp32.clamp(min=fp8_min, max=fp8_max).to( + torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn + ) + + B_fp32 = ( + (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max + ) + B = B_fp32.clamp(min=fp8_min, max=fp8_max).to( + torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn + ) + else: + int8_info = torch.iinfo(torch.int8) + int8_max, int8_min = int8_info.max, int8_info.min + + A_fp32 = ( + (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * int8_max + ) + A = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8) + + B_fp32 = ( + (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * int8_max + ) + B = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8) + + block_n, block_k = block_size[0], block_size[1] + n_tiles = (N + block_n - 1) // block_n + k_tiles = (K + block_k - 1) // block_k + + As = torch.rand(M, k_tiles, dtype=torch.float32, device="cuda") * factor_for_scale + Bs = ( + torch.rand(n_tiles, k_tiles, dtype=torch.float32, device="cuda") + * factor_for_scale + ) + + best_config = None + best_time = float("inf") + for config in tqdm(search_space): + try: + kernel_time = benchmark_config( + A, + B, + As, + Bs, + block_size, + config, + out_dtype, + num_iters=10, + ) + except triton.runtime.autotuner.OutOfResources: + # Some configurations may be invalid and fail to compile. + continue + + if kernel_time < best_time: + best_time = kernel_time + best_config = config + now = datetime.now() + print(f"{now.ctime()}] Completed tuning for batch_size={M}") + assert best_config is not None + return best_config + + +def save_configs( + N, + K, + block_n, + block_k, + configs, + save_path, + input_type="fp8", +) -> None: + os.makedirs(save_path, exist_ok=True) + device_name = get_device_name().replace(" ", "_") + json_file_name = f"N={N},K={K},device_name={device_name},dtype={input_type}_w8a8,block_shape=[{block_n}, {block_k}].json" + + config_file_path = os.path.join(save_path, json_file_name) + print(f"Writing best config to {config_file_path}...") + + with open(config_file_path, "w") as f: + json.dump(configs, f, indent=4) + f.write("\n") + + +def get_available_gpu_count(): + """Get the number of available GPUs.""" + return torch.cuda.device_count() + + +def tune_on_gpu(args_dict): + """Run tuning on a specific GPU.""" + gpu_id = args_dict["gpu_id"] + batch_sizes = args_dict["batch_sizes"] + weight_shapes = args_dict["weight_shapes"] + args = args_dict["args"] + + torch.cuda.set_device(gpu_id) + print(f"Starting tuning on GPU {gpu_id} with batch sizes {batch_sizes}") + + block_n = args.block_n + block_k = args.block_k + out_dtype = DTYPE_MAP[args.out_dtype] + save_path = args.save_path + input_type = args.input_type + + search_space = get_configs_compute_bound() + search_space = [ + config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0 + ] + + start = time.perf_counter() + results = {} + for shape in tqdm(weight_shapes, desc=f"GPU {gpu_id} - Shapes"): + N, K = shape[0], shape[1] + print(f"[GPU {gpu_id}] Tune for weight shape of `N: {N}, K: {K}`") + benchmark_results = [ + tune( + batch_size, + N, + K, + [block_n, block_k], + out_dtype, + search_space, + input_type, + ) + for batch_size in tqdm(batch_sizes, desc=f"GPU {gpu_id} - Batch sizes") + ] + best_configs = {M: config for M, config in zip(batch_sizes, benchmark_results)} + save_configs(N, K, block_n, block_k, best_configs, save_path, input_type) + + end = time.perf_counter() + print(f"Tuning on GPU {gpu_id} took {end - start:.2f} seconds") + + +def distribute_batch_sizes(batch_sizes, num_gpus): + """Distribute batch sizes across available GPUs.""" + batches_per_gpu = [] + for i in range(num_gpus): + start_idx = i * len(batch_sizes) // num_gpus + end_idx = (i + 1) * len(batch_sizes) // num_gpus + batches_per_gpu.append(batch_sizes[start_idx:end_idx]) + return batches_per_gpu + + +def main(args): + print(args) + + num_gpus = get_available_gpu_count() + if num_gpus == 0: + raise RuntimeError("No GPU available for tuning") + print(f"Found {num_gpus} GPUs for parallel tuning") + + torch.cuda.init() + + if args.batch_size is None: + batch_sizes = [ + 1, + 2, + 4, + 8, + 16, + 24, + 32, + 48, + 64, + 96, + 128, + 256, + 512, + 1024, + 1536, + 2048, + 3072, + 4096, + ] + else: + batch_sizes = [args.batch_size] + num_gpus = 1 # If only one batch size, use only one GPU + + weight_shapes = get_weight_shapes(args.tp_size) + + batches_per_gpu = distribute_batch_sizes(batch_sizes, num_gpus) + + process_args = [] + for gpu_id in range(num_gpus): + process_args.append( + { + "gpu_id": gpu_id, + "batch_sizes": batches_per_gpu[gpu_id], + "weight_shapes": weight_shapes, # Each GPU processes all weight shapes + "args": args, + } + ) + + ctx = mp.get_context("spawn") + with ctx.Pool(num_gpus) as pool: + pool.map(tune_on_gpu, process_args) + + print("Multi-GPU tuning completed") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument("--tp-size", "-tp", type=int, default=8) + parser.add_argument( + "--input-type", type=str, choices=["fp8", "int8"], default="fp8" + ) + parser.add_argument( + "--out-dtype", + type=str, + choices=["float32", "float16", "bfloat16", "half"], + default="float16", + ) + parser.add_argument("--block-n", type=int, default=128) + parser.add_argument("--block-k", type=int, default=128) + parser.add_argument("--batch-size", type=int, required=False) + parser.add_argument( + "--save-path", type=str, default="python/sglang/srt/layers/quantization/configs" + ) + args = parser.parse_args() + + main(args) diff --git a/benchmark/kernels/rmsnorm/benchmark_rmsnorm.py b/benchmark/kernels/rmsnorm/benchmark_rmsnorm.py new file mode 100644 index 000000000..aeeea62c0 --- /dev/null +++ b/benchmark/kernels/rmsnorm/benchmark_rmsnorm.py @@ -0,0 +1,230 @@ +import itertools +from typing import Optional, Tuple, Union + +import torch +import triton +from flashinfer.norm import fused_add_rmsnorm, rmsnorm +from torch import nn +from vllm import _custom_ops as vllm_ops + + +class HuggingFaceRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + orig_dtype = x.dtype + x = x.to(torch.float32) + if residual is not None: + x = x + residual.to(torch.float32) + residual = x.to(orig_dtype) + + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + self.variance_epsilon) + x = x.to(orig_dtype) * self.weight + if residual is None: + return x + else: + return x, residual + + +def rmsnorm_naive( + x: torch.Tensor, + weight: torch.Tensor, + residual: Optional[torch.Tensor] = None, + eps: float = 1e-6, +): + naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps) + naive_norm.weight = nn.Parameter(weight) + naive_norm = naive_norm.to(x.device) + + orig_shape = x.shape + x = x.view(-1, x.shape[-1]) + if residual is not None: + residual = residual.view(-1, residual.shape[-1]) + + output = naive_norm(x, residual) + + if isinstance(output, tuple): + output = (output[0].view(orig_shape), output[1].view(orig_shape)) + else: + output = output.view(orig_shape) + return output + + +def rmsnorm_flashinfer( + x: torch.Tensor, + weight: torch.Tensor, + residual: Optional[torch.Tensor] = None, + eps: float = 1e-6, +): + orig_shape = x.shape + x = x.view(-1, x.shape[-1]) + if residual is not None: + residual = residual.view(-1, residual.shape[-1]) + + if residual is not None: + fused_add_rmsnorm(x, residual, weight, eps) + output = (x, residual) + else: + output = rmsnorm(x, weight, eps) + + if isinstance(output, tuple): + output = (output[0].view(orig_shape), output[1].view(orig_shape)) + else: + output = output.view(orig_shape) + return output + + +def rmsnorm_vllm( + x: torch.Tensor, + weight: torch.Tensor, + residual: Optional[torch.Tensor] = None, + eps: float = 1e-6, +): + orig_shape = x.shape + x = x.view(-1, x.shape[-1]) + if residual is not None: + residual = residual.view(-1, residual.shape[-1]) + + if residual is not None: + vllm_ops.fused_add_rms_norm(x, residual, weight, eps) + output = (x, residual) + else: + out = torch.empty_like(x) + vllm_ops.rms_norm(out, x, weight, eps) + output = out + + if isinstance(output, tuple): + output = (output[0].view(orig_shape), output[1].view(orig_shape)) + else: + output = output.view(orig_shape) + return output + + +def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True): + dtype = torch.bfloat16 + x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda") + weight = torch.ones(hidden_size, dtype=dtype, device="cuda") + residual = torch.randn_like(x) if use_residual else None + + output_naive = rmsnorm_naive( + x.clone(), weight, residual.clone() if residual is not None else None + ) + output_flashinfer = rmsnorm_flashinfer( + x.clone(), weight, residual.clone() if residual is not None else None + ) + output_vllm = rmsnorm_vllm( + x.clone(), weight, residual.clone() if residual is not None else None + ) + + if use_residual: + output_naive = output_naive[0] + output_flashinfer = output_flashinfer[0] + output_vllm = output_vllm[0] + + print(f"Naive output={output_naive}") + print(f"FlashInfer output={output_flashinfer}") + print(f"VLLM output={output_vllm}") + + if torch.allclose( + output_naive, output_flashinfer, atol=1e-2, rtol=1e-2 + ) and torch.allclose(output_naive, output_vllm, atol=1e-2, rtol=1e-2): + print("✅ All implementations match") + else: + print("❌ Implementations differ") + + +batch_size_range = [2**i for i in range(0, 7, 2)] +seq_length_range = [2**i for i in range(6, 11, 1)] +head_num_range = [32, 48] +configs = list(itertools.product(head_num_range, batch_size_range, seq_length_range)) + + +def get_benchmark(use_residual): + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["head_num", "batch_size", "seq_len"], + x_vals=[list(_) for _ in configs], + line_arg="provider", + line_vals=["huggingface", "flashinfer", "vllm"], + line_names=["HuggingFace", "FlashInfer", "vLLM"], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="us", + plot_name=f"rmsnorm-performance-{'with' if use_residual else 'without'}-residual", + args={}, + ) + ) + def benchmark(head_num, batch_size, seq_len, provider): + dtype = torch.bfloat16 + hidden_size = head_num * 128 # assuming head_dim = 128 + + x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda") + weight = torch.ones(hidden_size, dtype=dtype, device="cuda") + residual = torch.randn_like(x) if use_residual else None + + quantiles = [0.5, 0.2, 0.8] + + if provider == "huggingface": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: rmsnorm_naive( + x.clone(), + weight, + residual.clone() if residual is not None else None, + ), + quantiles=quantiles, + ) + elif provider == "flashinfer": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: rmsnorm_flashinfer( + x.clone(), + weight, + residual.clone() if residual is not None else None, + ), + quantiles=quantiles, + ) + else: + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: rmsnorm_vllm( + x.clone(), + weight, + residual.clone() if residual is not None else None, + ), + quantiles=quantiles, + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--use_residual", action="store_true", help="Whether to use residual connection" + ) + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/rmsnorm/", + help="Path to save rmsnorm benchmark results", + ) + args = parser.parse_args() + + # Run correctness test + calculate_diff( + batch_size=4, seq_len=128, hidden_size=4096, use_residual=args.use_residual + ) + + # Get the benchmark function with proper use_residual setting + benchmark = get_benchmark(args.use_residual) + # Run performance benchmark + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py b/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py new file mode 100644 index 000000000..3e17205e7 --- /dev/null +++ b/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py @@ -0,0 +1,169 @@ +import os + +import torch +import triton +import triton.language as tl + + +@torch.compile(dynamic=True) +def get_last_loc_torch( + req_to_token: torch.Tensor, + req_pool_indices_tensor: torch.Tensor, + prefix_lens_tensor: torch.Tensor, +) -> torch.Tensor: + return torch.where( + prefix_lens_tensor > 0, + req_to_token[req_pool_indices_tensor, prefix_lens_tensor - 1], + torch.full_like(prefix_lens_tensor, -1), + ) + + +@triton.jit +def get_last_loc_kernel( + req_to_token, + req_pool_indices_tensor, + prefix_lens_tensor, + result, + num_tokens, + req_to_token_stride, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + mask = offset < num_tokens + + prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0) + req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0) + + token_mask = prefix_lens > 0 + token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1) + tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1) + + tl.store(result + offset, tokens, mask=mask) + + +def get_last_loc_triton( + req_to_token: torch.Tensor, + req_pool_indices_tensor: torch.Tensor, + prefix_lens_tensor: torch.Tensor, +) -> torch.Tensor: + BLOCK_SIZE = 256 + num_tokens = prefix_lens_tensor.shape[0] + result = torch.empty_like(prefix_lens_tensor) + grid = (triton.cdiv(num_tokens, BLOCK_SIZE),) + + get_last_loc_kernel[grid]( + req_to_token, + req_pool_indices_tensor, + prefix_lens_tensor, + result, + num_tokens, + req_to_token.stride(0), + BLOCK_SIZE, + ) + return result + + +def test_get_last_loc(): + max_batch = 4097 + max_context_len = 6148 + batch_size = 20 + + # Initialize input tensors + req_to_token = torch.zeros( + (max_batch, max_context_len), dtype=torch.int32, device="cuda" + ) + req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device="cuda") + pre_lens = torch.randint( + -max_context_len // 2, + max_context_len, + (batch_size,), + dtype=torch.int64, + device="cuda", + ) + + last_loc_res = get_last_loc_triton(req_to_token, req_pool_indices, pre_lens) + last_loc_ref = get_last_loc_torch(req_to_token, req_pool_indices, pre_lens) + + # Compare results + torch.testing.assert_close(last_loc_res, last_loc_ref) + + +def get_benchmark(): + batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=batch_sizes, + line_arg="provider", + line_vals=["reference", "triton"], + line_names=["PyTorch", "Triton"], + styles=[("blue", "-"), ("green", "-")], + ylabel="us", + plot_name="get-last-loc-performance", + args={}, + ) + ) + def benchmark(batch_size, provider): + max_batch = 2048 + max_context_len = 16384 + + req_to_token = torch.zeros( + (max_batch, max_context_len), dtype=torch.int32, device="cuda" + ) + req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device="cuda") + pre_lens = torch.randint( + -max_context_len // 2, + max_context_len, + (batch_size,), + dtype=torch.int64, + device="cuda", + ) + + quantiles = [0.5, 0.2, 0.8] + + if provider == "reference": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: get_last_loc_torch(req_to_token, req_pool_indices, pre_lens), + quantiles=quantiles, + ) + elif provider == "triton": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: get_last_loc_triton(req_to_token, req_pool_indices, pre_lens), + quantiles=quantiles, + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +def run_benchmark(save_path: str = "./configs/benchmark_ops/get_last_loc/"): + """Run benchmark and save results""" + + # Ensure save path exists + os.makedirs(save_path, exist_ok=True) + + # Run correctness test + test_get_last_loc() + print("Correctness test passed!") + + # Run performance test + benchmark = get_benchmark() + benchmark.run(print_data=True, save_path=save_path) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/get_last_loc/", + help="Path to save benchmark results", + ) + args = parser.parse_args() + + run_benchmark(args.save_path) diff --git a/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py b/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py new file mode 100644 index 000000000..1ce43c8ba --- /dev/null +++ b/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py @@ -0,0 +1,342 @@ +import itertools +import os + +import torch +import triton +import triton.language as tl + + +@triton.jit +def write_req_to_token_pool_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(0) + + req_pool_index = tl.load(req_pool_indices + pid) + pre_len = tl.load(pre_lens + pid) + seq_len = tl.load(seq_lens + pid) + + # TODO: optimize this? + cumsum_start = 0 + for i in range(pid): + cumsum_start += tl.load(extend_lens + i) + + num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < (seq_len - pre_len) + value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask) + tl.store( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + offset + + pre_len, + value, + mask=mask, + ) + + +@triton.jit +def write_req_to_token_pool_triton_optimize( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + req_to_token_ptr_stride: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_batch = tl.program_id(0) + pid_token = tl.program_id(1) + + req_pool_index = tl.load(req_pool_indices + pid_batch) + pre_len = tl.load(pre_lens + pid_batch) + seq_len = tl.load(seq_lens + pid_batch) + extend_len = seq_len - pre_len + + cumsum_start = 0 + for i in range(pid_batch): + cumsum_start += tl.load(extend_lens + i) + + token_start = pid_token * BLOCK_SIZE + + offset = tl.arange(0, BLOCK_SIZE) + actual_offset = token_start + offset + mask = actual_offset < extend_len + + src_ptr = out_cache_loc + cumsum_start + actual_offset + src_ptr = tl.max_contiguous(tl.multiple_of(src_ptr, BLOCK_SIZE), BLOCK_SIZE) + value = tl.load(src_ptr, mask=mask) + dst_ptr = ( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + actual_offset + + pre_len + ) + dst_ptr = tl.max_contiguous(tl.multiple_of(dst_ptr, BLOCK_SIZE), BLOCK_SIZE) + + tl.store(dst_ptr, value, mask=mask) + + +def write_req_to_token_pool_reference( + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + pre_lens: torch.Tensor, + seq_lens: torch.Tensor, + extend_lens: torch.Tensor, + out_cache_loc: torch.Tensor, +) -> None: + """Reference implementation using PyTorch""" + for i in range(len(req_pool_indices)): + req_pool_idx = req_pool_indices[i].item() + pre_len = pre_lens[i].item() + seq_len = seq_lens[i].item() + extend_len = extend_lens[i].item() + + cumsum_start = sum(extend_lens[:i].tolist()) + + # Copy values from out_cache_loc to req_to_token + req_to_token[req_pool_idx, pre_len:seq_len] = out_cache_loc[ + cumsum_start : cumsum_start + extend_len + ] + + +def test_write_req_to_token_pool(): + max_batch = 4097 + max_context_len = 6148 + batch_size = 1 + extend_len = 14 + + # Initialize input tensors + req_to_token = torch.zeros( + (max_batch, max_context_len), dtype=torch.int32, device="cuda" + ) + req_pool_indices = torch.tensor([42], dtype=torch.int32, device="cuda") + pre_lens = torch.tensor([8], dtype=torch.int32, device="cuda") + seq_lens = torch.tensor([22], dtype=torch.int32, device="cuda") + extend_lens = torch.tensor([extend_len], dtype=torch.int32, device="cuda") + out_cache_loc = torch.arange(extend_len, dtype=torch.int32, device="cuda") + + # Create copies for reference implementation + req_to_token_ref = req_to_token.clone() + req_to_token_opt = req_to_token.clone() + + # Run original triton kernel + write_req_to_token_pool_triton[(batch_size,)]( + req_to_token, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + ) + + # Run optimized triton kernel + def grid(batch_size, extend_len): + num_token_blocks = triton.cdiv(extend_len, 512) + return (batch_size, num_token_blocks) + + write_req_to_token_pool_triton_optimize[grid(batch_size, extend_len)]( + req_to_token_opt, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + BLOCK_SIZE=512, + ) + + # Run reference implementation + write_req_to_token_pool_reference( + req_to_token_ref, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + ) + + # Compare results + torch.testing.assert_close(req_to_token, req_to_token_ref) + torch.testing.assert_close(req_to_token_opt, req_to_token_ref) + + # Test case 2: batch size > 1 + batch_size = 3 + extend_lens_list = [14, 20, 30] + total_extend_len = sum(extend_lens_list) + + req_to_token = torch.zeros( + (max_batch, max_context_len), dtype=torch.int32, device="cuda" + ) + req_pool_indices = torch.tensor([42, 100, 200], dtype=torch.int32, device="cuda") + pre_lens = torch.tensor([8, 10, 15], dtype=torch.int32, device="cuda") + seq_lens = torch.tensor([22, 30, 45], dtype=torch.int32, device="cuda") + extend_lens = torch.tensor(extend_lens_list, dtype=torch.int32, device="cuda") + out_cache_loc = torch.arange(total_extend_len, dtype=torch.int32, device="cuda") + + req_to_token_ref = req_to_token.clone() + req_to_token_opt = req_to_token.clone() + + # Run original triton kernel + write_req_to_token_pool_triton[(batch_size,)]( + req_to_token, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + ) + + # Run optimized triton kernel + max_extend_len = max(extend_lens_list) + write_req_to_token_pool_triton_optimize[grid(batch_size, max_extend_len)]( + req_to_token_opt, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + BLOCK_SIZE=512, + ) + + # Run reference implementation + write_req_to_token_pool_reference( + req_to_token_ref, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + ) + + # Compare results + torch.testing.assert_close(req_to_token, req_to_token_ref) + torch.testing.assert_close(req_to_token_opt, req_to_token_ref) + + +def get_benchmark(): + batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128] + extend_lens = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] + configs = list(itertools.product(batch_sizes, extend_lens)) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "extend_len"], + x_vals=configs, + line_arg="provider", + line_vals=["reference", "triton", "triton_optimize"], + line_names=["PyTorch", "Triton", "Triton Optimized"], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="us", + plot_name="write-req-to-token-pool-performance", + args={}, + ) + ) + def benchmark(batch_size, extend_len, provider): + max_batch = 256 + max_context_len = 16384 + + extend_lens_list = [extend_len] * batch_size + total_extend_len = sum(extend_lens_list) + + req_to_token = torch.zeros( + (max_batch, max_context_len), dtype=torch.int32, device="cuda" + ) + req_pool_indices = torch.arange(batch_size, dtype=torch.int32, device="cuda") + pre_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda") * 8 + seq_lens = pre_lens + extend_len + extend_lens = torch.tensor(extend_lens_list, dtype=torch.int32, device="cuda") + out_cache_loc = torch.arange(total_extend_len, dtype=torch.int32, device="cuda") + + quantiles = [0.5, 0.2, 0.8] + + if provider == "reference": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: write_req_to_token_pool_reference( + req_to_token.clone(), + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + ), + quantiles=quantiles, + ) + elif provider == "triton": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: write_req_to_token_pool_triton[(batch_size,)]( + req_to_token.clone(), + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + ), + quantiles=quantiles, + ) + else: + + def run_optimized(): + block_size = 128 if extend_len <= 1024 else 512 + grid_config = (batch_size, triton.cdiv(extend_len, block_size)) + write_req_to_token_pool_triton_optimize[grid_config]( + req_to_token.clone(), + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + BLOCK_SIZE=block_size, + ) + + ms, min_ms, max_ms = triton.testing.do_bench( + run_optimized, quantiles=quantiles + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +def run_benchmark(save_path: str = "./configs/benchmark_ops/write_req_to_token_pool/"): + """Run benchmark and save results""" + + # Ensure save path exists + os.makedirs(save_path, exist_ok=True) + + # Run correctness test + test_write_req_to_token_pool() + print("Correctness test passed!") + + # Run performance test + benchmark = get_benchmark() + benchmark.run(print_data=True, save_path=save_path) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/write_req_to_token_pool/", + help="Path to save benchmark results", + ) + args = parser.parse_args() + + run_benchmark(args.save_path) diff --git a/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py b/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py new file mode 100644 index 000000000..98144d470 --- /dev/null +++ b/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py @@ -0,0 +1,283 @@ +import itertools + +import torch +import torch.nn.functional as F +import triton.testing as tt + +from sglang.srt.layers.attention.triton_ops.extend_attention import extend_attention_fwd + + +def extend_attention_fwd_torch( + q: torch.Tensor, # [extend_tokens, H_Q, D] + k: torch.Tensor, # [extend_tokens, H_KV, D] + v: torch.Tensor, # [extend_tokens, H_KV, D] + o: torch.Tensor, # [extend_tokens, H_Q, D] + k_cache: torch.Tensor, # [total_tokens, H_KV, D] + v_cache: torch.Tensor, # [total_tokens, H_KV, D] + qo_indptr: torch.Tensor, # [B+1] + kv_indptr: torch.Tensor, # [B+1] + kv_indices: torch.Tensor, # [prefix_tokens] + sliding_window_size: int, +): + B = qo_indptr.size(0) - 1 + _, H_Q, D = q.shape + _, H_KV, _ = k.shape + + group_size = H_Q // H_KV + scale = 1.0 / D**0.5 + + for i in range(B): + q_start = int(qo_indptr[i].item()) + q_end = int(qo_indptr[i + 1].item()) + kv_start = int(kv_indptr[i].item()) + kv_end = int(kv_indptr[i + 1].item()) + + prefix_indices = kv_indices[kv_start:kv_end] + k_prefix = k_cache[prefix_indices] # [prefix_len, H_KV, D] + v_prefix = v_cache[prefix_indices] # [prefix_len, H_KV, D] + + k_extend = k[q_start:q_end] # [extend_len, H_KV, D] + v_extend = v[q_start:q_end] # [extend_len, H_KV, D] + q_extend = q[q_start:q_end] # [extend_len, H_Q, D] + + k_full = torch.cat([k_prefix, k_extend], dim=0) # [total_len, H_KV, D] + v_full = torch.cat([v_prefix, v_extend], dim=0) # [total_len, H_KV, D] + + if group_size != 1: + k_full_hq = k_full.repeat_interleave( + group_size, dim=1 + ) # [total_len, H_Q, D] + v_full_hq = v_full.repeat_interleave( + group_size, dim=1 + ) # [total_len, H_Q, D] + else: + k_full_hq = k_full + v_full_hq = v_full + + prefix_len = k_prefix.size(0) + extend_len = k_extend.size(0) + total_len = prefix_len + extend_len + + # causal + pos_keys = torch.arange(total_len, device=q.device) + t = prefix_len + torch.arange(extend_len, device=q.device) # [extend_len] + causal_mask = pos_keys.unsqueeze(0) <= t.unsqueeze(1) + + # sliding window + if sliding_window_size is not None and sliding_window_size > 0: + start = (t - (sliding_window_size)).clamp_min(0) # [extend_len] + else: + start = torch.zeros_like(t) + window_mask = pos_keys.unsqueeze(0) >= start.unsqueeze(1) + + final_mask = causal_mask & window_mask + + attn_scores = ( + torch.einsum("qhd,khd->qhk", q_extend, k_full_hq) * scale + ) # [extend_len, H_Q, total_len] + attn_scores = attn_scores.masked_fill(~final_mask.unsqueeze(1), float("-inf")) + + attn_weights = F.softmax(attn_scores, dim=-1) + o[q_start:q_end] = torch.einsum("qhk,khd->qhd", attn_weights, v_full_hq) + + +def _build_batch( + B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE, dtype=torch.bfloat16, device="cuda" +): + b_seq_len_prefix = torch.randint( + 1, max(2, N_CTX // 2), (B,), dtype=torch.int32, device=device + ) + b_seq_len_extend = torch.randint( + 1, max(2, N_CTX // 2), (B,), dtype=torch.int32, device=device + ) + b_seq_len = b_seq_len_prefix + b_seq_len_extend + + b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device) + b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0) + b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device) + b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0) + + kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) + kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0) + + kv_indices = torch.zeros( + (int(b_seq_len_prefix.sum().item()),), dtype=torch.int32, device=device + ) + for i in range(B): + s = kv_indptr[i].item() + e = kv_indptr[i + 1].item() + kv_indices[s:e] = torch.arange( + b_start_loc[i], + b_start_loc[i] + b_seq_len_prefix[i], + dtype=torch.int32, + device=device, + ) + + total_token_num = int(torch.sum(b_seq_len).item()) + extend_token_num = int(torch.sum(b_seq_len_extend).item()) + + k_buffer = torch.empty( + (total_token_num, H_KV, D), dtype=dtype, device=device + ).normal_(mean=0.1, std=0.2) + v_buffer = torch.empty( + (total_token_num, H_KV, D), dtype=dtype, device=device + ).normal_(mean=0.1, std=0.2) + + k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) + + for i in range(B): + extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] + extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] + extend_start = b_start_loc_extend[i] + extend_end = b_start_loc_extend[i] + b_seq_len_extend[i] + + k_extend[extend_start:extend_end] = k_buffer[ + extend_start_in_buffer:extend_end_in_buffer + ] + v_extend[extend_start:extend_end] = v_buffer[ + extend_start_in_buffer:extend_end_in_buffer + ] + q_extend[extend_start:extend_end] = torch.empty( + (int(b_seq_len_extend[i].item()), H_Q, D), dtype=dtype, device=device + ).normal_(mean=0.1, std=0.2) + + o_extend_triton = torch.empty( + (extend_token_num, H_Q, D), dtype=dtype, device=device + ) + o_extend_torch = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) + + b_seq_len_extend = b_seq_len - b_seq_len_prefix + max_len_extend = int(torch.max(b_seq_len_extend, 0)[0].item()) + qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) + qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0) + + inputs = dict( + q_extend=q_extend, + k_extend=k_extend, + v_extend=v_extend, + k_buffer=k_buffer, + v_buffer=v_buffer, + o_extend_triton=o_extend_triton, + o_extend_torch=o_extend_torch, + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + max_len_extend=max_len_extend, + WINDOW_SIZE=WINDOW_SIZE, + ) + meta = dict( + B=B, N_CTX=N_CTX, H_Q=H_Q, H_KV=H_KV, D=D, extend_token_num=extend_token_num + ) + return inputs, meta + + +def _run_triton(inputs): + extend_attention_fwd( + inputs["q_extend"], + inputs["k_extend"], + inputs["v_extend"], + inputs["o_extend_triton"], + inputs["k_buffer"], + inputs["v_buffer"], + inputs["qo_indptr"], + inputs["kv_indptr"], + inputs["kv_indices"], + custom_mask=None, + is_causal=True, + mask_indptr=None, + max_len_extend=inputs["max_len_extend"], + sliding_window_size=inputs["WINDOW_SIZE"], + ) + + +def _run_torch_ref(inputs): + extend_attention_fwd_torch( + inputs["q_extend"], + inputs["k_extend"], + inputs["v_extend"], + inputs["o_extend_torch"], + inputs["k_buffer"], + inputs["v_buffer"], + inputs["qo_indptr"], + inputs["kv_indptr"], + inputs["kv_indices"], + inputs["WINDOW_SIZE"], + ) + + +N_CTXS = [1024, 2048, 4096, 8192] +WINDOW_SIZES = [-1, 127, 256, 512] + +CONFIGS = list(itertools.product(N_CTXS, WINDOW_SIZES)) + +PROVIDERS = ["torch", "triton"] + + +@tt.perf_report( + tt.Benchmark( + x_names=["N_CTX", "WINDOW_SIZE"], + x_vals=CONFIGS, + line_arg="provider", + line_vals=PROVIDERS, + line_names=PROVIDERS, + ylabel="Runtime (ms)", + plot_name="extend_attention_triton_vs_torch", + args={ + "B": 32, + "H_Q": 64, + "H_KV": 8, + "D": 128, + "dtype": "bf16", + "device": "cuda", + "check_correctness": False, + "warmup": 25, + "rep": 100, + }, + ) +) +def bench( + N_CTX, + provider, + B, + H_Q, + H_KV, + D, + dtype, + device, + WINDOW_SIZE, + check_correctness, + warmup, + rep, +): + torch.manual_seed(0) + torch.cuda.manual_seed(0) + dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} + dt = dtype_map[dtype] + + inputs, _ = _build_batch( + B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE, dtype=dt, device=device + ) + + if check_correctness and provider == "triton": + _run_triton(inputs) + _run_torch_ref(inputs) + torch.cuda.synchronize() + if not torch.allclose( + inputs["o_extend_triton"], inputs["o_extend_torch"], rtol=1e-3, atol=1e-3 + ): + raise AssertionError("Mismatch between triton and torch reference.") + + if provider == "triton": + ms = tt.do_bench(lambda: _run_triton(inputs), warmup=warmup, rep=rep) + elif provider == "torch": + ms = tt.do_bench(lambda: _run_torch_ref(inputs), warmup=warmup, rep=rep) + else: + raise ValueError(provider) + + return ms + + +if __name__ == "__main__": + bench.run(print_data=True, show_plots=False) diff --git a/benchmark/line_retrieval/README.md b/benchmark/line_retrieval/README.md new file mode 100644 index 000000000..3c81a63cb --- /dev/null +++ b/benchmark/line_retrieval/README.md @@ -0,0 +1,37 @@ +## Download data + +``` +wget https://raw.githubusercontent.com/merrymercy/merrymercy.github.io/master/files/random_words.json +python3 gen_data.py --number 1000 +``` + +## Run benchmark + +### Benchmark sglang +``` +python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-hf --port 30000 +``` + +``` +python3 bench_sglang.py --src-index 600 --num-q 50 --parallel 1 +``` + + +### + +``` +# original +Accuracy: 0.940, latency: 332.83 s + +# parallel encoding (no_adjust, offset = 1000) +Accuracy: 0.760, latency: 238.46 s + +# parallel encoding (no_adjust, offset = 3000) +Accuracy: 0.760, latency: 238.46 s + +# parallel encoding (no_adjust, offset = 0) +Accuracy: 0.520, latency: 238.46 s + +# parallel encoding (adjust_cache) +Accuracy: 0.460, latency: 257.66 s +``` diff --git a/benchmark/line_retrieval/bench_sglang.py b/benchmark/line_retrieval/bench_sglang.py new file mode 100644 index 000000000..e974e7dd3 --- /dev/null +++ b/benchmark/line_retrieval/bench_sglang.py @@ -0,0 +1,149 @@ +import argparse +import json +import re +import time + +import numpy as np + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +@sgl.function +def line_retrieval(s, prefix, suffix, body_0, body_1, body_2, body_3): + s += prefix + "\n" + + contexts = [body_0, body_1, body_2, body_3] + position_ids_offset = [i * 1000 for i in range(len(contexts))] + forks = s.fork(len(contexts), position_ids_offset) + forks += lambda i: contexts[i] + "\n" + forks.join(mode="concate_and_append") + + s += "\n" + suffix + s += sgl.gen("answer", max_tokens=16) + + +def eval_model(args, line_obj, num_hoops, src_indices, dst_percents): + arguments = [] + labels = [] + sum_src_indices = [] + sum_dst_indices = [] + + for i in range(len(src_indices)): + for j in range(len(dst_percents)): + src_index = src_indices[i] + dst_percent = dst_percents[j] + + query_indices = line_obj["group_by_num_hoops"][str(num_hoops)] + query_indices = [ + q + for q in query_indices + if all(l <= src_index for l in line_obj["links"][q]) and q < src_index + ] + dst_index = query_indices[ + min(int(len(query_indices) * dst_percent), len(query_indices) - 1) + ] + label = line_obj["values"][dst_index] + + body = line_obj["lines"][: src_index + 1] + suffix = line_obj["suffix"].replace("???", line_obj["indices"][dst_index]) + body_part_len = len(body) // 4 + + arguments.append( + { + "prefix": line_obj["prefix"], + "body_0": "\n".join(body[:body_part_len]), + "body_1": "\n".join(body[body_part_len : 2 * body_part_len]), + "body_2": "\n".join(body[2 * body_part_len : 3 * body_part_len]), + "body_3": "\n".join(body[3 * body_part_len :]), + "suffix": suffix, + } + ) + labels.append(label) + sum_src_indices.append(src_index) + sum_dst_indices.append(dst_index) + + # Select backend + backend = select_sglang_backend(args) + + tic = time.perf_counter() + states = line_retrieval.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + corrects = [] + for i in range(len(arguments)): + output = states[i]["answer"] + prompt_len = states[i].get_meta_info("answer").get("prompt_length", -1) + label = labels[i] + + # Try all numbers + findall = re.findall("\d+", output) + if not findall: + response_number = output + else: + for response_number in findall: + if response_number == label: + break + + correct = response_number == label + corrects.append(correct) + + # Log results + summary = ( + f"Line index: {sum_src_indices[i]} -> {sum_dst_indices[i]}, " + f"Prompt len: {prompt_len}, " + f"Correct: {correct}, " + f"Label: {label}, Predicted: {response_number}, " + ) + print(summary) + + accuracy = np.mean(corrects) + print(f"Accuracy: {accuracy:.3f}, latency: {latency:.2f} s") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "line_retrieval", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": len(arguments), + "other": { + "num_questions": len(arguments), + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +def main(args): + line_obj = json.load(open(args.data_path, "r")) + + num_hoops = args.num_hoops + for src_index in args.src_index: + src_indices = [src_index] + num_queries = args.num_queries_per_src + dst_percents = [i * (1 / (num_queries)) for i in range(num_queries)] + eval_model(args, line_obj, num_hoops, src_indices, dst_percents) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="lines_1000_0.0.json") + parser.add_argument("--src-index", type=int, nargs="+", default=[100]) + parser.add_argument("--num-queries-per-src", type=int, default=10) + parser.add_argument("--num-hoops", type=int, default=1) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/line_retrieval/gen_data.py b/benchmark/line_retrieval/gen_data.py new file mode 100644 index 000000000..a01d40a06 --- /dev/null +++ b/benchmark/line_retrieval/gen_data.py @@ -0,0 +1,139 @@ +""" +Generate line data for line retrieval task. + +Usage: +python3 gen_data.py --number 1000 +""" + +import argparse +import json +from collections import defaultdict + +import numpy as np +from tqdm import tqdm + + +def generate_lines(random_words, num_lines, redirect_ratio): + prefix = "Here is a list of lines, each with its corresponding REGISTER_CONTENT value. Please memorize them. Be prepared to provide the REGISTER_CONTENT value for a specific line index when I ask." + suffix = "The list has ended. Please give the final REGISTER_CONTENT value for a specific line after resolving the redirections and references. For example, the REGISTER_CONTENT of Line __idx0__ is __val0__. The REGISTER_CONTENT of Line __idx1__ is __val1__. The REGISTER_CONTENT of Line __idx2__ is __val2__. The REGISTER_CONTENT of Line ??? is" + + # Raw lines + visited_indices = set([None]) + visited_values = set([None]) + + lines = [] + redirects = [] + indices = [] + values = [] + for i in tqdm(range(num_lines)): + line_index = None + while line_index in visited_indices: + line_index = "-".join(np.random.choice(random_words, size=(2,))) + visited_indices.add(line_index) + + line_value = np.random.randint(low=0, high=999999) + line_value = f"{line_value:06}" + + line = f"Line {line_index}: The REGISTER_CONTENT is {line_value}." + lines.append(line) + redirects.append(None) + indices.append(line_index) + values.append(line_value) + + # Add redirect + if redirect_ratio > 0: + num_redirect_lines = int(len(lines) * redirect_ratio) + redirect_indices = np.random.choice( + np.arange(len(lines)), size=(num_redirect_lines,), replace=False + ) + for i in redirect_indices: + target_idx = np.random.choice(min(i * 2 + 100, num_lines)) + lines[i] = ( + f"Line {indices[i]}: The REGISTER_CONTENT is the same as Line {indices[target_idx]}." + ) + redirects[i] = target_idx + + # Build links and find sources + links = [[] for _ in range(num_lines)] + contains_ring = set() + for i in range(num_lines): + if redirects[i] is None: + continue + + tmp_link = [] + cur = i + visited = set() + while redirects[cur] is not None: + visited.add(cur) + tmp_link.append(redirects[cur]) + cur = redirects[cur] + + if cur in visited: + contains_ring.add(i) + tmp_link = None + break + values[i] = values[cur] + links[i] = tmp_link + + # Group by num_links + group_by_num_hoops = defaultdict(list) + for i in range(num_lines): + if i in contains_ring: + continue + group_by_num_hoops[len(links[i]) + 1].append(i) + + keys = sorted(list(group_by_num_hoops.keys())) + for num_links in keys: + print(f"#links: {num_links}, #lines: {len(group_by_num_hoops[num_links])}") + + # Append few-shot examples + hoop1_candidates = list(group_by_num_hoops[1]) + hoop1_candidate_keys = {c: max([c] + links[c]) for c in hoop1_candidates} + hoop1_candidates.sort(key=lambda c: hoop1_candidate_keys[c]) + hoop2_candidates = list(group_by_num_hoops[2]) + hoop2_candidate_keys = {c: max([c] + links[c]) for c in hoop2_candidates} + hoop2_candidates.sort(key=lambda c: hoop2_candidate_keys[c]) + + i = hoop1_candidates[5] + suffix = suffix.replace("__idx0__", indices[i]).replace("__val0__", values[i]) + if len(hoop2_candidates): + i = hoop2_candidates[0] + suffix = suffix.replace("__idx1__", indices[i]).replace("__val1__", values[i]) + i = hoop2_candidates[1] + suffix = suffix.replace("__idx2__", indices[i]).replace("__val2__", values[i]) + else: + i = hoop1_candidates[1] + suffix = suffix.replace("__idx1__", indices[i]).replace("__val1__", values[i]) + i = hoop1_candidates[10] + suffix = suffix.replace("__idx2__", indices[i]).replace("__val2__", values[i]) + + obj = { + "prefix": prefix, + "suffix": suffix, + "lines": lines, + "indices": indices, + "values": values, + "links": links, + "group_by_num_hoops": group_by_num_hoops, + "contains_ring": sorted(list(contains_ring)), + } + return obj + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--number", type=int) + parser.add_argument("--redirect-ratio", type=float, default=0.0) + args = parser.parse_args() + + num_lines = args.number + + random_words_filename = "random_words.json" + random_words = json.load(open(random_words_filename, "r")) + + np.random.seed(42) + obj = generate_lines(random_words, num_lines, args.redirect_ratio) + + fout = f"lines_{num_lines}_{args.redirect_ratio:.1f}.json" + with open(fout, "w") as fout: + json.dump(obj, fout, indent=2) diff --git a/benchmark/llava_bench/README.md b/benchmark/llava_bench/README.md new file mode 100644 index 000000000..a71f339df --- /dev/null +++ b/benchmark/llava_bench/README.md @@ -0,0 +1,61 @@ +## Download benchmark images + +``` +python3 download_images.py +``` + +image benchmark source: https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild + +### Other Dependency +``` +pip3 install "sglang[all]" +pip3 install "torch>=2.1.2" "transformers>=4.36" pillow +``` + +## Run benchmark + +### Benchmark sglang +Launch a server +``` +python3 -m sglang.launch_server --model-path liuhaotian/llava-v1.6-vicuna-7b --tokenizer-path llava-hf/llava-1.5-7b-hf --port 30000 +``` + +Run benchmark +``` +# Run with local models +python3 bench_sglang.py --num-questions 60 + +# Run with OpenAI models +python3 bench_sglang.py --num-questions 60 --backend gpt-4-vision-preview +``` + +### Bench LLaVA original code +``` +git clone git@github.com:haotian-liu/LLaVA.git +cd LLaVA +git reset --hard 9a26bd1435b4ac42c282757f2c16d34226575e96 +pip3 install -e . + +cd ~/sglang/benchmark/llava_bench +CUDA_VISIBLE_DEVICES=0 bash bench_hf_llava_bench.sh +``` + + +### Benchmark llama.cpp + +``` +# Install +CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python +pip install sse_starlette starlette_context pydantic_settings + +# Download weights +mkdir -p ~/model_weights/llava-v1.5-7b/ +wget https://huggingface.co/mys/ggml_llava-v1.5-7b/resolve/main/ggml-model-f16.gguf -O ~/model_weights/llava-v1.5-7b/ggml-model-f16.gguf +wget https://huggingface.co/mys/ggml_llava-v1.5-7b/resolve/main/mmproj-model-f16.gguf -O ~/model_weights/llava-v1.5-7b/mmproj-model-f16.gguf +``` + +``` +python3 -m llama_cpp.server --model ~/model_weights/llava-v1.5-7b/ggml-model-f16.gguf --clip_model_path ~/model_weights/llava-v1.5-7b/mmproj-model-f16.gguf --chat_format llava-1-5 --port 23000 + +OPENAI_BASE_URL=http://localhost:23000/v1 python3 bench_sglang.py --backend gpt-4-vision-preview --num-q 1 +``` diff --git a/benchmark/llava_bench/bench_hf_llava_bench.sh b/benchmark/llava_bench/bench_hf_llava_bench.sh new file mode 100755 index 000000000..a51a715b6 --- /dev/null +++ b/benchmark/llava_bench/bench_hf_llava_bench.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +python -m llava.eval.model_vqa \ + --model-path liuhaotian/llava-v1.5-7b \ + --question-file ./questions.jsonl \ + --image-folder ./images \ + --answers-file ./answers_hf.jsonl \ + --temperature 0 \ + --conv-mode vicuna_v1 diff --git a/benchmark/llava_bench/bench_hf_mme.sh b/benchmark/llava_bench/bench_hf_mme.sh new file mode 100755 index 000000000..6ed332fe0 --- /dev/null +++ b/benchmark/llava_bench/bench_hf_mme.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +python -m llava.eval.model_vqa_loader \ + --model-path liuhaotian/llava-v1.5-7b \ + --question-file ./mme_pack/llava_mme_bench_replace.jsonl \ + --image-folder ./mme_pack/MME_Benchmark_release_version \ + --answers-file ./answers_hf_mme.jsonl \ + --temperature 0 \ + --conv-mode vicuna_v1 diff --git a/benchmark/llava_bench/bench_sglang.py b/benchmark/llava_bench/bench_sglang.py new file mode 100644 index 000000000..b9e8c1405 --- /dev/null +++ b/benchmark/llava_bench/bench_sglang.py @@ -0,0 +1,96 @@ +import argparse +import json +import os +import time + +import tqdm + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +@sgl.function +def image_qa(s, image_file, question): + s += sgl.user(sgl.image(image_file) + question) + s += sgl.assistant(sgl.gen("answer", max_tokens=args.max_tokens)) + + +def main(args): + lines = list(read_jsonl(args.question_file))[: args.num_questions] + arguments = [ + { + "image_file": os.path.abspath(args.image_folder + "/" + l["image"]), + "question": l["text"], + } + for l in lines + ] + # arguments = [ + # {"image_file": + # Image.open(os.path.abspath(args.image_folder + "/" + l["image"])), + # "question": l["text"]} for l in lines + # ] + + states = [None] * len(lines) + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm.tqdm(range(len(lines))): + image_file = arguments[i]["image_file"] + question = arguments[i]["question"] + ret = image_qa.run(image_file=image_file, question=question, temperature=0) + states[i] = ret + else: + states = image_qa.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.perf_counter() - tic + + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + print(f"Write output to {args.answer_file}") + with open(args.answer_file, "w") as fout: + for i in range(len(lines)): + value = { + "question_id": lines[i]["question_id"], + "prompt": lines[i]["text"], + "text": states[i]["answer"].strip(), + "model_id": backend.model_info["model_path"], + "answer_id": i, + "metadata": {}, + } + fout.write(json.dumps(value) + "\n") + + with open(args.result_file, "a") as fout: + value = { + "task": "llava_bench", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": len(lines), + "parallel": args.parallel, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--question-file", type=str, default="questions.jsonl") + parser.add_argument("--answer-file", type=str, default="answers.jsonl") + parser.add_argument("--image-folder", type=str, default="./images") + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--num-questions", type=int, default=None) + parser.add_argument("--max-tokens", type=int, default=768) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/llava_bench/bench_sglang_mme.sh b/benchmark/llava_bench/bench_sglang_mme.sh new file mode 100644 index 000000000..8c8c09cc0 --- /dev/null +++ b/benchmark/llava_bench/bench_sglang_mme.sh @@ -0,0 +1,2 @@ +MME_FOLDER=./mme_pack +python3 bench_sglang.py --num-questions 5000 --question-file $MME_FOLDER/llava_mme_bench_replace.jsonl --answer-file answer_mme.jsonl --image-folder $MME_FOLDER/MME_Benchmark_release_version --max-tokens 4 diff --git a/benchmark/llava_bench/download_images.py b/benchmark/llava_bench/download_images.py new file mode 100644 index 000000000..701190a03 --- /dev/null +++ b/benchmark/llava_bench/download_images.py @@ -0,0 +1,20 @@ +import os + +# Create the 'images' directory if it doesn't exist +if not os.path.exists("images"): + os.makedirs("images") + +# Base URL +base_url = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/" + +# Loop through image numbers +for i in range(1, 25): + # Format the image number with leading zeros + image_number = str(i).zfill(3) + image_url = base_url + image_number + ".jpg" + image_path = "images/" + image_number + ".jpg" + + # Download the image using wget + os.system(f"wget -O {image_path} {image_url}") + +print("Download complete.") diff --git a/benchmark/llm_judge/README.md b/benchmark/llm_judge/README.md new file mode 100644 index 000000000..164b19a08 --- /dev/null +++ b/benchmark/llm_judge/README.md @@ -0,0 +1,33 @@ +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 25 --parallel 8 +python3 bench_sglang.py --num-questions 16 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --backend vllm --num-questions 25 +``` + + +### Benchmark guidance +``` +python3 bench_other.py --backend guidance --num-questions 25 --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --backend lmql --num-questions 25 --parallel 1 +``` diff --git a/benchmark/llm_judge/bench_other.py b/benchmark/llm_judge/bench_other.py new file mode 100644 index 000000000..8e6029067 --- /dev/null +++ b/benchmark/llm_judge/bench_other.py @@ -0,0 +1,151 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +system_prompt = "Please serve as an impartial judge and rigorously evaluate the quality of the following article. Apply the most stringent standards possible, showing no leniency." + +dimension_prompts = [ + "Content: This refers to the essences of the essay. The substance should be well researched, accurate, relevant to the topic and should show a thorough understanding of the subject. The essay should also reflect a clear goal or purpose.", + "Organization and Structure: An essay needs to be properly structured with a clear introduction, body, and conclusion. The essay should flow naturally, with one paragraph leading seamlessly into the next.", + "Argument and Analysis: The argument made in the essay should be logical, coherent and clearly articulated. Each point made should be backed up by solid evidence and thorough analysis.", + "Clarity and Precision: The essay should be written in a clear and concise manner. The points made should be easily understood by the reader. The language used should also be precise and unambiguous.", + "Grammar and Punctuation: Proper use of grammar and punctuation is vital in an academic essay. Errors in grammar and punctuation not only distract the reader but can also negatively impact the meaning and interpretation of the content.", + "Referencing and Citation: An essay should contain proper citations and references for all sources used. This not only prevents accusations of plagiarism but also gives credit to the authors of the works that have contributed to the essay. The citation should adhere to a specific format as required by the academic institution or specified by the professor.", +] + + +def multi_dimension_judge(article, generate): + s = system_prompt + s += "\n```\n" + article + "\n```\n\n" + + judges = [] + for i in range(len(dimension_prompts)): + comp = generate( + s + + "USER: Please judge the quality based on the following metric. " + + dimension_prompts[i] + + " Please provide a single-paragraph judgement. " + + "Focus on the provided metric and do not say other things. " + 'End your judgement paragraph with the word "END"\nJUDGE:', + max_tokens=256, + stop="END", + ) + judges.append(comp) + + s += "I will judge the quality based on the following metrics.\n" + for i in range(len(dimension_prompts)): + s += dimension_prompts[i].split(":")[0] + ": " + judges[i].strip() + "\n" + + s += "In summary, on a scale of 1 to 10, I would give the article a score of" + s += generate(s, max_tokens=2, stop=None) + + return s + + +async def multi_dimension_judge_async(article, generate): + s = system_prompt + s += "\n```\n" + article + "\n```\n\n" + + judges = [] + for i in range(len(dimension_prompts)): + comp = await generate( + s + + "USER: Please judge the quality based on the following metric. " + + dimension_prompts[i] + + " Please provide a single-paragraph judgement. " + + "Focus on the provided metric and do not say other things. " + 'End your judgement paragraph with the word "END"\nJUDGE:', + max_tokens=256, + stop="END", + ) + judges.append(comp) + + s += "I will judge the quality based on the following metrics.\n" + for i in range(len(dimension_prompts)): + s += dimension_prompts[i].split(":")[0] + ": " + judges[i].strip() + "\n" + + s += "In summary, on a scale of 1 to 10, I would give the article a score of" + s += await generate(s, max_tokens=2, stop=None) + + return s + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + states = [None] * len(lines) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + tic = time.perf_counter() + + if args.backend != "lmql": + + def get_one_answer(i): + states[i] = multi_dimension_judge(lines[i], call_generate) + + if args.parallel == 1: + for i in tqdm(range(len(lines))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(lines)))), + total=len(lines), + ) + ) + + else: + import asyncio + + async def get_one_answer_async(i): + states[i] = await multi_dimension_judge_async(lines[i], call_generate) + + batches = [] + for i in range(0, len(lines), args.parallel): + batches.append(list(range(i, min(i + args.parallel, len(lines))))) + + loop = asyncio.get_event_loop() + for bt in tqdm(batches): + loop.run_until_complete( + asyncio.gather(*[get_one_answer_async(i) for i in bt]) + ) + + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "llm_judge", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="articles.jsonl") + parser.add_argument("--num-questions", type=int, default=20) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/llm_judge/bench_sglang.py b/benchmark/llm_judge/bench_sglang.py new file mode 100644 index 000000000..97e6c3979 --- /dev/null +++ b/benchmark/llm_judge/bench_sglang.py @@ -0,0 +1,97 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +system_prompt = "Please serve as an impartial judge and rigorously evaluate the quality of the following article. Apply the most stringent standards possible, showing no leniency." + +dimension_prompts = [ + "Content: This refers to the essences of the essay. The substance should be well researched, accurate, relevant to the topic and should show a thorough understanding of the subject. The essay should also reflect a clear goal or purpose.", + "Organization and Structure: An essay needs to be properly structured with a clear introduction, body, and conclusion. The essay should flow naturally, with one paragraph leading seamlessly into the next.", + "Argument and Analysis: The argument made in the essay should be logical, coherent and clearly articulated. Each point made should be backed up by solid evidence and thorough analysis.", + "Clarity and Precision: The essay should be written in a clear and concise manner. The points made should be easily understood by the reader. The language used should also be precise and unambiguous.", + "Grammar and Punctuation: Proper use of grammar and punctuation is vital in an academic essay. Errors in grammar and punctuation not only distract the reader but can also negatively impact the meaning and interpretation of the content.", + "Referencing and Citation: An essay should contain proper citations and references for all sources used. This not only prevents accusations of plagiarism but also gives credit to the authors of the works that have contributed to the essay. The citation should adhere to a specific format as required by the academic institution or specified by the professor.", +] + + +@sgl.function +def multi_dimension_judge(s, article): + s += system_prompt + s += "\n```\n" + article + "\n```\n\n" + + forks = s.fork(len(dimension_prompts)) + for i in range(len(dimension_prompts)): + forks[i] += ( + "USER: Please judge the quality based on the following metric. " + + dimension_prompts[i] + + " Please provide a single-paragraph judgement. " + + "Focus on the provided metric and do not say other things. " + 'End your judgement paragraph with the word "END"\nJUDGE:' + ) + forks[i] += sgl.gen("judgement", max_tokens=256, stop="END") + forks.join() + + s += "I will judge the quality based on the following metrics.\n" + for i in range(len(dimension_prompts)): + s += ( + dimension_prompts[i].split(":")[0] + + ": " + + forks[i]["judgement"].strip() + + "\n" + ) + + s += "In summary, on a scale of 1 to 10, I would give the article a score of" + s += sgl.gen("score", max_tokens=2) + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + arguments = [{"article": l} for l in lines] + + # Select backend + backend = select_sglang_backend(args) + + # Run requests + tic = time.perf_counter() + states = multi_dimension_judge.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "llm_judge", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="articles.jsonl") + parser.add_argument("--num-questions", type=int, default=20) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/long_json_decode/README.md b/benchmark/long_json_decode/README.md new file mode 100644 index 000000000..37fceee13 --- /dev/null +++ b/benchmark/long_json_decode/README.md @@ -0,0 +1,33 @@ +## Run benchmark + +### Benchmark sglang +``` +python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-instruct-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 5 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model codellama/CodeLlama-7b-instruct-hf --disable-log-requests --port 21000 --gpu 0.97 +``` + +``` +python3 bench_other.py --backend vllm --num-questions 5 +``` + + +### Benchmark guidance +``` +python3 bench_other.py --backend guidance --num-questions 5 --parallel 1 --n-ctx 11000 --model-path path/to/code-llama/gguf +``` + + +### Build dataset +``` +pip install wikipedia +python3 build_dataset.py +``` diff --git a/benchmark/long_json_decode/bench_other.py b/benchmark/long_json_decode/bench_other.py new file mode 100644 index 000000000..0ad38a014 --- /dev/null +++ b/benchmark/long_json_decode/bench_other.py @@ -0,0 +1,89 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + + +def json_decode(document, generate): + s = "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += "{\n" + s += ' "name": "' + s += generate(s, max_tokens=8, stop='"') + '",\n' + s += ' "country": "' + s += generate(s, max_tokens=8, stop='"') + '",\n' + s += ' "air port code": "' + s += generate(s, max_tokens=8, stop='"') + '",\n' + s += ' "top 3 landmarks": "' + s += generate(s, max_tokens=24, stop='"') + '",\n' + s += "}\n" + return s + + +def main(args): + lines = read_jsonl(args.data_path) + arguments = [] + for i in range(len(lines[: args.num_questions])): + arguments.append( + { + "document": lines[i]["document"], + } + ) + states = [None] * len(arguments) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + def get_one_answer(i): + states[i] = json_decode(generate=call_generate, **arguments[i]) + + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(arguments))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(arguments)))), + total=len(arguments), + ) + ) + + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "long_json_decode", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/long_json_decode/bench_sglang.py b/benchmark/long_json_decode/bench_sglang.py new file mode 100644 index 000000000..8394cfc2e --- /dev/null +++ b/benchmark/long_json_decode/bench_sglang.py @@ -0,0 +1,81 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +@sgl.function +def json_decode(s, document): + s += "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += "{\n" + s += ' "name": "' + sgl.gen("name", max_tokens=8, stop='"') + '",\n' + s += ' "country": "' + sgl.gen("country", max_tokens=8, stop='"') + '",\n' + s += ( + ' "air port code": "' + + sgl.gen("air port code", max_tokens=8, stop='"') + + '",\n' + ) + s += ( + ' "top 3 landmarks": "' + + sgl.gen("landmarks", max_tokens=24, stop='"') + + '",\n' + ) + s += "}\n" + + +def main(args): + lines = read_jsonl(args.data_path) + arguments = [] + for i in range(len(lines[: args.num_questions])): + arguments.append( + { + "document": lines[i]["document"], + } + ) + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.perf_counter() + states = json_decode.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "long_json_decode", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=10) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/long_json_decode/build_dataset.py b/benchmark/long_json_decode/build_dataset.py new file mode 100644 index 000000000..7b15d3951 --- /dev/null +++ b/benchmark/long_json_decode/build_dataset.py @@ -0,0 +1,27 @@ +import json + +import transformers +import wikipedia + +name = "meta-llama/Llama-2-7b-chat-hf" +t = transformers.AutoTokenizer.from_pretrained(name) +city_names = ["los angles", "london", "tokyo", "beijing", "singapore"] + + +for city_name in city_names: + content = str(wikipedia.page(city_name).content) + content = content.replace("\n\n", "\n") + + tokens = t.encode(content) + + truncate_len = int((10000 / len(tokens)) * len(content)) + truncate_content = content[:truncate_len] + truncate_tokens = t.encode(truncate_content) + + # Count token + print( + f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}" + ) + + with open("questions.jsonl", "a") as fout: + fout.write(json.dumps({"document": truncate_content}) + "\n") diff --git a/benchmark/lora/launch_server.py b/benchmark/lora/launch_server.py new file mode 100644 index 000000000..b0781ca30 --- /dev/null +++ b/benchmark/lora/launch_server.py @@ -0,0 +1,75 @@ +import argparse +import os + +NUM_LORAS = 4 +LORA_PATH = { + "base": "meta-llama/Llama-2-7b-hf", + "lora": "winddude/wizardLM-LlaMA-LoRA-7B", +} + + +def launch_server(args): + base_path = LORA_PATH["base"] + lora_path = LORA_PATH["lora"] + + if args.base_only: + cmd = f"python3 -m sglang.launch_server --model {base_path} " + else: + cmd = f"python3 -m sglang.launch_server --model {base_path} --lora-paths " + for i in range(NUM_LORAS): + lora_name = f"lora{i}" + cmd += f"{lora_name}={lora_path} " + cmd += f"--disable-radix " + cmd += f"--max-loras-per-batch {args.max_loras_per_batch} " + cmd += f"--max-running-requests {args.max_running_requests} " + cmd += f"--lora-backend {args.lora_backend} " + cmd += f"--tp-size {args.tp_size} " + if args.disable_custom_all_reduce: + cmd += "--disable-custom-all-reduce" + if args.enable_mscclpp: + cmd += "--enable-mscclpp" + print(cmd) + os.system(cmd) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--base-only", + action="store_true", + ) + parser.add_argument( + "--max-loras-per-batch", + type=int, + default=8, + ) + parser.add_argument( + "--max-running-requests", + type=int, + default=8, + ) + parser.add_argument( + "--lora-backend", + type=str, + default="triton", + ) + parser.add_argument( + "--tp-size", + type=int, + default=1, + help="Tensor parallel size for distributed inference", + ) + # disable_custom_all_reduce + parser.add_argument( + "--disable-custom-all-reduce", + action="store_true", + help="Disable custom all reduce when device does not support p2p communication", + ) + parser.add_argument( + "--enable-mscclpp", + action="store_true", + help="Enable using mscclpp for small messages for all-reduce kernel and fall back to NCCL.", + ) + args = parser.parse_args() + + launch_server(args) diff --git a/benchmark/lora/lora_bench.py b/benchmark/lora/lora_bench.py new file mode 100644 index 000000000..0a1e37a5c --- /dev/null +++ b/benchmark/lora/lora_bench.py @@ -0,0 +1,476 @@ +# Copyright 2023-2024 SGLang 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 asyncio +import json +import random +import resource +import sys +import time +import traceback +from argparse import ArgumentParser +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + +import aiohttp +import numpy as np +from launch_server import LORA_PATH, NUM_LORAS +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +from sglang.bench_serving import ( + AIOHTTP_TIMEOUT, + RequestFuncInput, + RequestFuncOutput, + calculate_metrics, + get_request, + get_tokenizer, + remove_prefix, + sample_random_requests, +) + +global args + + +# set ignore_eos True by default +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + # assert api_url.endswith( + # "completions" + # ), "OpenAI Completions API URL must end with 'completions'." + + prompt = request_func_input.prompt + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + # payload = { + # "model": request_func_input.model, + # "prompt": prompt, + # "temperature": 0.0, + # "best_of": 1, + # "max_tokens": request_func_input.output_len, + # "stream": not args.disable_stream, + # "ignore_eos": not args.disable_ignore_eos, + # **request_func_input.extra_request_body, + # } + # headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} + if args.base_only: + payload = { + "text": prompt, + "sampling_params": {"max_new_tokens": request_func_input.output_len}, + } + else: + payload = { + "text": prompt, + "sampling_params": {"max_new_tokens": request_func_input.output_len}, + "lora_path": f"lora{random.randint(0, NUM_LORAS - 1)}", + } + headers = {"Authorization": ""} + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["text"]: + # if data["choices"][0]["text"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + # generated_text += data["choices"][0]["text"] + generated_text += data["text"] + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = request_func_input.output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +ASYNC_REQUEST_FUNCS = { + "sglang": async_request_openai_completions, +} + + +async def benchmark( + backend: str, + api_url: str, + model_id: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: List[Tuple[str, int, int]], + request_rate: float, + disable_tqdm: bool, + extra_request_body: Dict[str, Any], +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + print("Starting initial single prompt test run...") + test_request = input_requests[0] + test_input = RequestFuncInput( + model=model_id, + prompt=test_request.prompt, + api_url=api_url, + prompt_len=test_request.prompt_len, + output_len=test_request.output_len, + lora_name="dummy", # the lora_name argument will not be used + image_data=None, + extra_request_body=extra_request_body, + ) + test_output = await request_func(request_func_input=test_input) + if not test_output.success: + raise ValueError( + "Initial test run failed - Please make sure benchmark arguments " + f"are correctly specified. Error: {test_output.error}" + ) + else: + print("Initial test run completed. Starting main benchmark run...") + + pbar = None if disable_tqdm else tqdm(total=len(input_requests)) + + benchmark_start_time = time.perf_counter() + tasks: List[asyncio.Task] = [] + async for request in get_request(input_requests, request_rate): + request_func_input = RequestFuncInput( + model=model_id, + prompt=request.prompt, + api_url=api_url, + prompt_len=request.prompt_len, + output_len=request.output_len, + lora_name="dummy", + image_data=None, + extra_request_body=extra_request_body, + ) + tasks.append( + asyncio.create_task( + request_func(request_func_input=request_func_input, pbar=pbar) + ) + ) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + if pbar is not None: + pbar.close() + + benchmark_duration = time.perf_counter() - benchmark_start_time + + metrics, output_lens = calculate_metrics( + input_requests=input_requests, + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + backend=backend, + ) + + print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) + print("{:<40} {:<10}".format("Backend:", backend)) + print("{:<40} {:<10}".format("Traffic request rate:", request_rate)) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) + print( + "{:<40} {:<10}".format( + "Total generated tokens (retokenized):", metrics.total_output_retokenized + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", metrics.request_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", metrics.input_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", metrics.output_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format("Total throughput (tok/s):", metrics.total_throughput) + ) + print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) + print( + "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) + ) + print( + "{:<40} {:<10.2f}".format( + "Median E2E Latency (ms):", metrics.median_e2e_latency_ms + ) + ) + print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) + print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) + print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) + print( + "{s:{c}^{n}}".format(s="Time per Output Token (excl. 1st token)", n=50, c="-") + ) + print("{:<40} {:<10.2f}".format("Mean TPOT (ms):", metrics.mean_tpot_ms)) + print("{:<40} {:<10.2f}".format("Median TPOT (ms):", metrics.median_tpot_ms)) + print("{:<40} {:<10.2f}".format("P99 TPOT (ms):", metrics.p99_tpot_ms)) + print("{s:{c}^{n}}".format(s="Inter-token Latency", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) + print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) + print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) + print("=" * 50) + + if ( + metrics.median_ttft_ms is not None + and metrics.mean_itl_ms is not None + and metrics.output_throughput is not None + ): + result = { + "backend": args.backend, + "request_rate": request_rate, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "median_itl_ms": metrics.median_itl_ms, + "output_throughput": metrics.output_throughput, + "random_input_len": args.random_input_len, + "random_output_len": args.random_output_len, + "random_range_ratio": args.random_range_ratio, + "duration": benchmark_duration, + "completed": metrics.completed, + } + else: + print(f"Error running benchmark for request rate: {request_rate}") + print("-" * 30) + + # Determine output file name + if args.output_file: + output_file_name = args.output_file + else: + now = datetime.now().strftime("%m%d") + output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" + + # Append results to a JSONL file + with open(output_file_name, "a") as file: + file.write(json.dumps(result) + "\n") + + result = { + "duration": benchmark_duration, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "mean_ttft_ms": metrics.mean_ttft_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "std_ttft_ms": metrics.std_ttft_ms, + "p99_ttft_ms": metrics.p99_ttft_ms, + "mean_tpot_ms": metrics.mean_tpot_ms, + "median_tpot_ms": metrics.median_tpot_ms, + "std_tpot_ms": metrics.std_tpot_ms, + "p99_tpot_ms": metrics.p99_tpot_ms, + "mean_itl_ms": metrics.mean_itl_ms, + "median_itl_ms": metrics.median_itl_ms, + "std_itl_ms": metrics.std_itl_ms, + "p99_itl_ms": metrics.p99_itl_ms, + "input_lens": [output.prompt_len for output in outputs], + "output_lens": output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + } + return result + + +def run_benchmark(args_: argparse.Namespace): + global args + args = args_ + + # Set global environments + set_ulimit() + random.seed(args.seed) + np.random.seed(args.seed) + + # Set url + if args.port is None: + args.port = { + "sglang": 30000, + }.get(args.backend, 30000) + + # api_url = ( + # f"{args.base_url}/v1/completions" + # if args.base_url + # else f"http://{args.host}:{args.port}/v1/completions" + # ) + api_url = ( + f"{args.base_url}/generate" + if args.base_url + else f"http://{args.host}:{args.port}/generate" + ) + + print(f"{args}\n") + + # Read dataset + backend = args.backend + model_id = args.model = LORA_PATH["base"] + tokenizer_id = args.model + + tokenizer = get_tokenizer(tokenizer_id) + + input_requests = sample_random_requests( + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + dataset_path="", + ) + + return asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + model_id=model_id, + tokenizer=tokenizer, + input_requests=input_requests, + request_rate=args.request_rate, + disable_tqdm=False, + extra_request_body={}, + ) + ) + + +def set_ulimit(target_soft_limit=65535): + resource_type = resource.RLIMIT_NOFILE + current_soft, current_hard = resource.getrlimit(resource_type) + + if current_soft < target_soft_limit: + try: + resource.setrlimit(resource_type, (target_soft_limit, current_hard)) + except ValueError as e: + print(f"Fail to set RLIMIT_NOFILE: {e}") + + +if __name__ == "__main__": + parser = ArgumentParser(description="Benchmark the online lora serving throughput.") + parser.add_argument( + "--backend", + type=str, + choices=list(ASYNC_REQUEST_FUNCS.keys()), + default="sglang", + help="Must specify a backend, depending on the LLM Inference Engine.", + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + parser.add_argument( + "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." + ) + parser.add_argument( + "--port", + type=int, + help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=50, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=1024, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + type=int, + default=128, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=0.0, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " + "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", + ) + parser.add_argument( + "--base-only", + action="store_true", + ) + parser.add_argument("--output-file", type=str, help="Output JSONL file name.") + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + args = parser.parse_args() + run_benchmark(args) diff --git a/benchmark/mmlu/README.md b/benchmark/mmlu/README.md new file mode 100644 index 000000000..16de20cda --- /dev/null +++ b/benchmark/mmlu/README.md @@ -0,0 +1,59 @@ +## Download data +``` +bash download_data.sh +``` + +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --nsub 10 +``` + +``` +# OpenAI models +python3 bench_sglang.py --backend gpt-3.5-turbo --parallel 8 +``` + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --nsub 10 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 + +# V100 +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 4500 --port 22000 +``` + +``` +python3 bench_other.py --nsub 10 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --nsub 10 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + + +### Benchmark lmql +``` +CUDA_VISIBLE_DEVICES=0,1 lmql serve-model meta-llama/Llama-2-7b-chat-hf --cuda --port 23000 +``` + +``` +python3 bench_other.py --nsub 10 --backend lmql --parallel 2 +``` diff --git a/benchmark/mmlu/bench_other.py b/benchmark/mmlu/bench_other.py new file mode 100644 index 000000000..f1b166c2b --- /dev/null +++ b/benchmark/mmlu/bench_other.py @@ -0,0 +1,173 @@ +import argparse +import asyncio +import json +import os +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import pandas as pd +import tiktoken +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate + +choices = ["A", "B", "C", "D"] + +tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo") + + +def format_subject(subject): + l = subject.split("_") + s = "" + for entry in l: + s += " " + entry + return s + + +def format_example(df, idx, include_answer=True): + prompt = df.iloc[idx, 0] + k = df.shape[1] - 2 + for j in range(k): + prompt += "\n{}. {}".format(choices[j], df.iloc[idx, j + 1]) + prompt += "\nAnswer:" + if include_answer: + prompt += " {}\n\n".format(df.iloc[idx, k + 1]) + return prompt + + +def gen_prompt(train_df, subject, k=-1): + prompt = "The following are multiple choice questions (with answers) about{}.\n\n".format( + format_subject(subject) + ) + if k == -1: + k = train_df.shape[0] + for i in range(k): + prompt += format_example(train_df, i) + return prompt + + +def evaluate(args, subject, dev_df, test_df, call_generate): + prompts = [] + labels = [] + + # Construct prompts + k = args.ntrain + train_prompt = gen_prompt(dev_df, subject, k) + while len(tokenizer.encode(train_prompt)) > 1536: + k -= 1 + train_prompt = gen_prompt(dev_df, subject, k) + + for i in range(test_df.shape[0]): + prompt_end = format_example(test_df, i, include_answer=False) + prompt = train_prompt + prompt_end + prompts.append(prompt) + + label = test_df.iloc[i, test_df.shape[1] - 1] + labels.append(label) + + preds = [None] * len(prompts) + max_tokens = 1 + + # Run requests + if args.backend != "lmql": + # Use thread pool + def get_one_answer(i): + pred = call_generate(prompts[i], temperature=0, max_tokens=max_tokens) + preds[i] = pred.strip()[0] + + tic = time.perf_counter() + if args.parallel == 1: + for i in range(len(prompts)): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + executor.map(get_one_answer, list(range(len(prompts)))) + else: + # Use asyncio + async def batched_call(batch_size): + for i in range(0, len(prompts), batch_size): + tasks = [] + for p in prompts[i : i + batch_size]: + tasks.append(call_generate(p, temperature=0, max_tokens=max_tokens)) + rets = await asyncio.gather(*tasks) + for j in range(len(rets)): + preds[i + j] = rets[j].strip()[0] + + tic = time.perf_counter() + asyncio.run(batched_call(batch_size=args.parallel)) + latency = time.perf_counter() - tic + + # Compute accuracy + cors = [pred == label for pred, label in zip(preds, labels)] + acc = np.mean(cors) + cors = np.array(cors) + + print( + "Average accuracy {:.3f}, latency {:.2f}, #q: {} - {}".format( + acc, latency, len(prompts), subject + ) + ) + + return cors, acc, latency + + +def main(args): + subjects = sorted( + [ + f.split("_test.csv")[0] + for f in os.listdir(os.path.join(args.data_dir, "test")) + if "_test.csv" in f + ] + ) + + all_cors = [] + all_latencies = [] + num_requests = 0 + + # Select backend + call_generate = get_call_generate(args) + + for subject in tqdm(subjects[: args.nsub]): + dev_df = pd.read_csv( + os.path.join(args.data_dir, "dev", subject + "_dev.csv"), header=None + )[: args.ntrain] + test_df = pd.read_csv( + os.path.join(args.data_dir, "test", subject + "_test.csv"), header=None + ) + + cors, acc, latency = evaluate(args, subject, dev_df, test_df, call_generate) + all_cors.append(cors) + all_latencies.append(latency) + num_requests += len(test_df) + + total_latency = np.sum(all_latencies) + print("Total latency: {:.3f}".format(total_latency)) + + weighted_acc = np.mean(np.concatenate(all_cors)) + print("Average accuracy: {:.3f}".format(weighted_acc)) + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "mmlu", + "backend": args.backend, + "num_gpus": 1, + "latency": round(total_latency, 3), + "accuracy": round(weighted_acc, 3), + "num_requests": num_requests, + "other": { + "nsub": args.nsub, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--ntrain", type=int, default=5) + parser.add_argument("--data_dir", type=str, default="data") + parser.add_argument("--nsub", type=int, default=60) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/mmlu/bench_sglang.py b/benchmark/mmlu/bench_sglang.py new file mode 100644 index 000000000..23057be4a --- /dev/null +++ b/benchmark/mmlu/bench_sglang.py @@ -0,0 +1,181 @@ +import argparse +import json +import os +import time + +import numpy as np +import pandas as pd +import tiktoken + +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + dump_bench_raw_result, + select_sglang_backend, +) + +choices = ["A", "B", "C", "D"] + +tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo") + + +def format_subject(subject): + l = subject.split("_") + s = "" + for entry in l: + s += " " + entry + return s + + +def format_example(df, idx, include_answer=True): + prompt = df.iloc[idx, 0] + k = df.shape[1] - 2 + for j in range(k): + prompt += "\n{}. {}".format(choices[j], df.iloc[idx, j + 1]) + prompt += "\nAnswer:" + if include_answer: + prompt += " {}\n\n".format(df.iloc[idx, k + 1]) + return prompt + + +def gen_prompt(train_df, subject, k=-1): + prompt = "The following are multiple choice questions (with answers) about{}.\n\n".format( + format_subject(subject) + ) + if k == -1: + k = train_df.shape[0] + for i in range(k): + prompt += format_example(train_df, i) + return prompt + + +def main(args): + subjects = sorted( + [ + f.split("_test.csv")[0] + for f in os.listdir(os.path.join(args.data_dir, "test")) + if "_test.csv" in f + ] + ) + + # Build prompts + arguments = [] + labels = [] + num_questions = [] + + for subject in subjects[: args.nsub]: + dev_df = pd.read_csv( + os.path.join(args.data_dir, "dev", subject + "_dev.csv"), header=None + )[: args.ntrain] + test_df = pd.read_csv( + os.path.join(args.data_dir, "test", subject + "_test.csv"), header=None + ) + num_questions.append(test_df.shape[0]) + + k = args.ntrain + few_shot_examples = gen_prompt(dev_df, subject, k) + while len(tokenizer.encode(few_shot_examples)) > 1536: + k -= 1 + few_shot_examples = gen_prompt(dev_df, subject, k) + + for i in range(test_df.shape[0]): + prompt_end = format_example(test_df, i, include_answer=False) + + arguments.append( + { + "examples": few_shot_examples, + "question": prompt_end, + } + ) + + label = test_df.iloc[i, test_df.shape[1] - 1] + labels.append(label) + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + if args.backend.startswith("gpt-"): + + @sgl.function + def few_shot_mmlu(s, examples, question): + s += sgl.user(examples + question) + s += sgl.assistant(sgl.gen("answer")) + + else: + + @sgl.function + def few_shot_mmlu(s, examples, question): + s += examples + question + sgl.gen("answer") + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Select backend + backend = select_sglang_backend(args) + + # Run + tic = time.perf_counter() + states = few_shot_mmlu.run_batch( + arguments, + temperature=0, + max_new_tokens=1, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + preds = [ + s["answer"].strip()[0] if len(s["answer"].strip()) > 0 else "" for s in states + ] + latency = time.perf_counter() - tic + + # Compute accuracy + cors = [pred == label for pred, label in zip(preds, labels)] + + pt = 0 + for subject, num_qs in zip(subjects[: args.nsub], num_questions): + print( + f"subject: {subject}, #q:{num_qs}, acc: {np.mean(cors[pt: pt + num_qs]):.3f}" + ) + pt += num_qs + assert pt == len(cors) + weighted_acc = np.mean(cors) + + dump_bench_raw_result( + path=args.raw_result_file, + states=states, + preds=preds, + labels=labels, + ) + + # Print results + print("Total latency: {:.3f}".format(latency)) + print("Average accuracy: {:.3f}".format(weighted_acc)) + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "mmlu", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(weighted_acc, 3), + "num_requests": len(arguments), + "other": { + "nsub": args.nsub, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--ntrain", "-k", type=int, default=5) + parser.add_argument("--data_dir", "-d", type=str, default="data") + parser.add_argument("--save_dir", "-s", type=str, default="results") + parser.add_argument("--nsub", type=int, default=60) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/mmlu/download_data.sh b/benchmark/mmlu/download_data.sh new file mode 100644 index 000000000..fb224165e --- /dev/null +++ b/benchmark/mmlu/download_data.sh @@ -0,0 +1,2 @@ +wget https://people.eecs.berkeley.edu/~hendrycks/data.tar +tar xf data.tar diff --git a/benchmark/mmmu/README.md b/benchmark/mmmu/README.md new file mode 100644 index 000000000..80db21921 --- /dev/null +++ b/benchmark/mmmu/README.md @@ -0,0 +1,46 @@ +## Run evaluation + +### Evaluate sglang + +Host the VLM: + +``` +python -m sglang.launch_server --model-path Qwen/Qwen2-VL-7B-Instruct --port 30000 +``` + +It's recommended to reduce the memory usage by appending something like `--mem-fraction-static 0.6` to the command above. + +Benchmark: + +``` +python benchmark/mmmu/bench_sglang.py --port 30000 --concurrency 16 +``` + +You can adjust the `--concurrency` to control the number of concurrent OpenAI calls. + +You can use `--lora-path` to specify the LoRA adapter to apply during benchmarking. E.g., +``` +# Launch server with LoRA enabled +python -m sglang.launch_server --model-path microsoft/Phi-4-multimodal-instruct --port 30000 --trust-remote-code --disable-radix-cache --lora-paths vision= + +# Apply LoRA adapter during inferencing +python -m benchmark/mmmu/bench_sglang.py --concurrency 8 --lora-path vision +``` + +You can use `--response-answer-regex` to specify how to extract the answer from the response string. E.g., +``` +python3 -m sglang.launch_server --model-path zai-org/GLM-4.1V-9B-Thinking --reasoning-parser glm45 + +python3 bench_sglang.py --response-answer-regex "<\|begin_of_box\|>(.*)<\|end_of_box\|>" --concurrency 64 +``` + +You can use `--extra-request-body` to specify additional OpenAI request parameters. E.g., +``` +python3 bench_sglang.py --extra-request-body '{"max_new_tokens": 128, "temperature": 0.01}' +``` + +### Evaluate hf + +``` +python benchmark/mmmu/bench_hf.py --model-path Qwen/Qwen2-VL-7B-Instruct +``` diff --git a/benchmark/mmmu/bench_hf.py b/benchmark/mmmu/bench_hf.py new file mode 100644 index 000000000..949b63b80 --- /dev/null +++ b/benchmark/mmmu/bench_hf.py @@ -0,0 +1,164 @@ +import argparse + +import PIL +import torch +from data_utils import save_json +from eval_utils import ( + EvalArgs, + eval_result, + get_sampling_params, + prepare_samples, + process_result, +) +from tqdm import tqdm +from transformers import AutoModel, AutoProcessor, GenerationConfig + + +@torch.no_grad() +def eval_mmmu(args): + eval_args = EvalArgs.from_cli_args(args) + + sampling_params = get_sampling_params(eval_args) + generation_config = GenerationConfig( + max_new_tokens=sampling_params["max_new_tokens"], + do_sample=False, + ) + + try: + from transformers import AutoModelForImageTextToText + + model = AutoModelForImageTextToText.from_pretrained( + args.model_path, + torch_dtype="auto", + trust_remote_code=True, + ) + except Exception as first_exception: + try: + # check if the model is belongs to internvl + if "InternVL" in args.model_path: + from internvl_utils import load_image + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.model_path) + model = AutoModel.from_pretrained( + args.model_path, + torch_dtype="auto", + trust_remote_code=True, + ) + generation_config_internvl = dict( + max_new_tokens=sampling_params["max_new_tokens"], do_sample=False + ) + + else: + model = AutoModel.from_pretrained( + args.model_path, + torch_dtype="auto", + trust_remote_code=True, + init_tts=False, + ) + except Exception as second_exception: + raise RuntimeError( + f"Failed to load model: First attempt failed with {first_exception}, " + f"second attempt failed with {second_exception}" + ) from second_exception + + model = model.eval().cuda() + + processor = AutoProcessor.from_pretrained( + args.model_path, torch_dtype="auto", device_map="auto", trust_remote_code=True + ) + + samples = prepare_samples(eval_args) + out_samples = dict() + + answer_dict = {} + for sample in tqdm(samples): + prompt = sample["final_input_prompt"] + image = sample["image"] + prefix = prompt.split("<")[0] + suffix = prompt.split(">")[1] + assert image is not None + + if "InternVL" in args.model_path: + pixel_values = load_image(sample["image_path"]).to(torch.bfloat16).cuda() + contents = "" + if prefix: + contents += prefix + contents += "\n" + if suffix: + contents += suffix + response = model.chat( + tokenizer, pixel_values, contents, generation_config_internvl + ) + print(f"response: {response}") + process_result(response, sample, answer_dict, out_samples) + continue + + contents = [] + if prefix: + contents += [{"type": "text", "text": prefix}] + contents += [ + { + "type": "image", + "image": sample["image_path"], + } + ] + if suffix: + contents += [{"type": "text", "text": suffix}] + messages = [{"role": "user", "content": contents}] + try: + model_inputs = processor.tokenizer.apply_chat_template( + messages, + tokenize=True, + return_dict=True, + add_generation_prompt=True, + return_tensors="pt", + ).to(model.device) + input_len = model_inputs["input_ids"].shape[-1] + generation = model.generate( + **model_inputs, generation_config=generation_config + ) + generation = generation[0][input_len:] + response = processor.decode(generation, skip_special_tokens=True) + except: + contents = [] + if prefix: + contents += [prefix] + image = PIL.Image.open(sample["image_path"]) + contents += [image] + if suffix: + contents += [suffix] + messages = [{"role": "user", "content": contents}] + response = model.chat( + msgs=messages, + tokenizer=processor.tokenizer, + sampling=False, + max_new_tokens=sampling_params["max_new_tokens"], + use_tts_template=False, + generate_audio=False, + temperature=0.0, + ) + print(f"response: {response}") + process_result(response, sample, answer_dict, out_samples) + + args.output_path = f"{args.model_path}_answer_hf.json" + save_json(args.output_path, out_samples) + eval_result( + model_answer_path=args.output_path, + answer_dict=answer_dict, + eval_output_path=f"{args.model_path}_val_hf.json", + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--model-path", + type=str, + help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.", + required=True, + ) + EvalArgs.add_cli_args(parser) + args = parser.parse_args() + + eval_mmmu(args) diff --git a/benchmark/mmmu/bench_sglang.py b/benchmark/mmmu/bench_sglang.py new file mode 100644 index 000000000..d8834ea5f --- /dev/null +++ b/benchmark/mmmu/bench_sglang.py @@ -0,0 +1,212 @@ +""" +Bench the sglang-hosted vLM with benchmark MMMU + +Usage: + Host the VLM: python -m sglang.launch_server --model-path Qwen/Qwen2-VL-7B-Instruct --port 30000 + + Benchmark: python benchmark/mmmu/bench_sglang.py --port 30000 --concurrency 16 + +The eval output will be logged +""" + +import argparse +import asyncio +import re +import sys +import time +import traceback +from dataclasses import dataclass, field +from typing import Any, List, Optional, Tuple + +import aiohttp +import openai +from data_utils import save_json +from eval_utils import ( + EvalArgs, + eval_result, + get_sampling_params, + prepare_samples, + process_result, +) +from tqdm import tqdm + +from sglang.test.test_utils import add_common_sglang_args_and_parse + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) + + +@dataclass +class RequestFuncOutput: + generated_text: List[str] = field(default_factory=list) + prompt_len: List[int] = field(default_factory=list) + output_len: List[int] = field(default_factory=list) + latency: List[float] = field(default_factory=list) + ttft: List[float] = field(default_factory=list) + itl: List[float] = field(default_factory=list) # List of inter-token latencies + + success: bool = False + error: str = "" + + +async def async_request_profile(api_url: str) -> RequestFuncOutput: + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + output = RequestFuncOutput() + try: + async with session.post(url=api_url) as response: + if response.status == 200: + output.success = True + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + return output + + +def _get_prefix_suffix(prompt: str) -> Tuple[str, str]: + """Split the prompt into prefix and suffix.""" + prefix = prompt.split("<")[0] + suffix = prompt.split(">", 1)[1] + return prefix, suffix + + +async def process_sample( + client: Any, sample: dict, sampling_params: dict, lora_path: Optional[str] = None +) -> Tuple[dict, str]: + """Send a single sample to the LLM and return (sample, response).""" + prompt = sample["final_input_prompt"] + prefix, suffix = _get_prefix_suffix(prompt) + image = sample["image"] + assert image is not None + image_path = sample["image_path"] + extra_body = None if lora_path is None else {"lora_path": lora_path} + response = await client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": prefix}, + {"type": "image_url", "image_url": {"url": image_path}}, + {"type": "text", "text": suffix}, + ], + } + ], + temperature=0, + max_completion_tokens=sampling_params["max_new_tokens"], + max_tokens=sampling_params["max_new_tokens"], + extra_body=extra_body, + ) + return sample, response.choices[0].message.content + + +async def process_sample_with_semaphore( + semaphore: asyncio.Semaphore, + client: Any, + sample: dict, + sampling_params: dict, + lora_path: Optional[str] = None, +) -> Tuple[dict, str]: + """Wrap process_sample with a semaphore for concurrency control.""" + async with semaphore: + return await process_sample(client, sample, sampling_params, lora_path) + + +async def eval_mmmu(args) -> None: + """Main evaluation loop with concurrency control.""" + eval_args = EvalArgs.from_cli_args(args) + sampling_params = get_sampling_params(eval_args) + samples = prepare_samples(eval_args) + lora_path = eval_args.lora_path + answer_dict = {} + out_samples = {} + client = openai.AsyncOpenAI( + api_key="sk", base_url=f"http://127.0.0.1:{args.port}/v1" + ) + start = time.perf_counter() + base_url = f"http://127.0.0.1:{args.port}" + + if args.profile: + print("Starting profiler...") + profile_output = await async_request_profile( + api_url=f"{base_url}/start_profile" + ) + if profile_output.success: + print("Profiler started") + + samples = samples[: args.profile_number] + + if args.concurrency == 1: + # For concurrency == 1, run in sequential mode to ensure consistent order + # this is mainly for profiling + for sample in tqdm(samples): + _, response = await process_sample( + client, sample, sampling_params, lora_path + ) + answer = ( + re.search(args.response_answer_regex, response) + if response is not None + else None + ) + process_result( + answer.group(1) if answer else response, + sample, + answer_dict, + out_samples, + ) + else: + semaphore = asyncio.Semaphore(args.concurrency) + tasks = [ + process_sample_with_semaphore( + semaphore, client, sample, sampling_params, lora_path + ) + for sample in samples + ] + + for coro in tqdm(asyncio.as_completed(tasks), total=len(tasks)): + sample, response = await coro + answer = ( + re.search(args.response_answer_regex, response) + if response is not None + else None + ) + process_result( + answer.group(1) if answer else response, + sample, + answer_dict, + out_samples, + ) + + if args.profile: + print("Stopping profiler...") + profile_output = await async_request_profile(api_url=f"{base_url}/stop_profile") + if profile_output.success: + print("Profiler stopped") + + print(f"Benchmark time: {time.perf_counter() - start}") + args.output_path = "./answer_sglang.json" + save_json(args.output_path, out_samples) + eval_result( + model_answer_path=args.output_path, + answer_dict=answer_dict, + eval_output_path="./val_sglang.json", + ) + + +def parse_args(): + parser = argparse.ArgumentParser() + EvalArgs.add_cli_args(parser) + args = add_common_sglang_args_and_parse(parser) + return args + + +def main(): + args = parse_args() + asyncio.run(eval_mmmu(args)) + + +if __name__ == "__main__": + main() diff --git a/benchmark/mmmu/data_utils.py b/benchmark/mmmu/data_utils.py new file mode 100644 index 000000000..cf8916934 --- /dev/null +++ b/benchmark/mmmu/data_utils.py @@ -0,0 +1,221 @@ +"""Utils for data load, save, and process (e.g., prompt construction)""" + +import json +import os +import re + +import yaml + +DOMAIN_CAT2SUB_CAT = { + "Art and Design": ["Art", "Art_Theory", "Design", "Music"], + "Business": ["Accounting", "Economics", "Finance", "Manage", "Marketing"], + "Science": [ + "Biology", + "Chemistry", + "Geography", + "Math", + "Physics", + ], + "Health and Medicine": [ + "Basic_Medical_Science", + "Clinical_Medicine", + "Diagnostics_and_Laboratory_Medicine", + "Pharmacy", + "Public_Health", + ], + "Humanities and Social Science": [ + "History", + "Literature", + "Sociology", + "Psychology", + ], + "Tech and Engineering": [ + "Agriculture", + "Architecture_and_Engineering", + "Computer_Science", + "Electronics", + "Energy_and_Power", + "Materials", + "Mechanical_Engineering", + ], +} + + +CAT_SHORT2LONG = { + "acc": "Accounting", + "agri": "Agriculture", + "arch": "Architecture_and_Engineering", + "art": "Art", + "art_theory": "Art_Theory", + "bas_med": "Basic_Medical_Science", + "bio": "Biology", + "chem": "Chemistry", + "cli_med": "Clinical_Medicine", + "cs": "Computer_Science", + "design": "Design", + "diag_med": "Diagnostics_and_Laboratory_Medicine", + "econ": "Economics", + "elec": "Electronics", + "ep": "Energy_and_Power", + "fin": "Finance", + "geo": "Geography", + "his": "History", + "liter": "Literature", + "manage": "Manage", + "mark": "Marketing", + "mate": "Materials", + "math": "Math", + "mech": "Mechanical_Engineering", + "music": "Music", + "phar": "Pharmacy", + "phys": "Physics", + "psy": "Psychology", + "pub_health": "Public_Health", + "socio": "Sociology", +} + + +# DATA SAVING +def save_json(filename, ds): + with open(filename, "w") as f: + json.dump(ds, f, indent=4) + + +def get_multi_choice_info(options): + """ + Given the list of options for multiple choice question + Return the index2ans and all_choices + """ + + start_chr = "A" + all_choices = [] + index2ans = {} + for i, option in enumerate(options): + index2ans[chr(ord(start_chr) + i)] = option + all_choices.append(chr(ord(start_chr) + i)) + + return index2ans, all_choices + + +def load_yaml(file_path): + with open(file_path, "r") as stream: + try: + yaml_dict = yaml.safe_load(stream) + except yaml.YAMLError as exc: + print(exc) + + return yaml_dict + + +def parse_img_path(text): + matches = re.findall("", text) + return matches + + +def process_single_sample(data): + question = data["question"] + o_imgs_paths = [] + for option in data["options"]: + current_o_imgs_paths = parse_img_path(option) + for img_path in current_o_imgs_paths: + o_imgs_paths.append(img_path) + + if len(o_imgs_paths) > 1: # multiple images in options, used for random selection + return { + "id": data["id"], + "question": question, + "options": data["options"], + "answer": data["answer"], + "image": None, + "question_type": data["question_type"], + } + else: + return { + "id": data["id"], + "question": question, + "options": data["options"], + "answer": data["answer"], + "image": data["image_1"], + "question_type": data["question_type"], + } + + +# DATA SAVING +def save_json(filename, ds): + print(f"answers saved to: {filename}") + os.makedirs(os.path.dirname(filename), exist_ok=True) + with open(filename, "w") as f: + json.dump(ds, f, indent=4) + + +def save_jsonl(filename, data): + """ + Save a dictionary of data to a JSON Lines file with the filename as key and caption as value. + + Args: + filename (str): The path to the file where the data should be saved. + data (dict): The dictionary containing the data to save where key is the image path and value is the caption. + """ + with open(filename, "w", encoding="utf-8") as f: + for img_path, caption in data.items(): + # Extract the base filename without the extension + base_filename = os.path.basename(img_path) + # Create a JSON object with the filename as the key and caption as the value + json_record = json.dumps({base_filename: caption}, ensure_ascii=False) + # Write the JSON object to the file, one per line + f.write(json_record + "\n") + + +def save_args(args, path_dir): + argsDict = args.__dict__ + with open(path_dir + "setting.txt", "w") as f: + f.writelines("------------------ start ------------------" + "\n") + for eachArg, value in argsDict.items(): + f.writelines(eachArg + " : " + str(value) + "\n") + f.writelines("------------------- end -------------------") + + +# DATA PROCESSING +def construct_prompt(sample, config): + question = sample["question"] + options = eval(sample["options"]) + example = "" + if sample["question_type"] == "multiple-choice": + start_chr = "A" + prediction_range = [] + index2ans = {} + for option in options: + prediction_range.append(start_chr) + example += f"({start_chr}) {option}\n" + index2ans[start_chr] = option + start_chr = chr(ord(start_chr) + 1) + empty_prompt_sample_structure = config["multi_choice_example_format"] + empty_prompt = empty_prompt_sample_structure.format(question, example) + res_dict = {} + res_dict["index2ans"] = index2ans + res_dict["correct_choice"] = sample["answer"] + res_dict["all_choices"] = prediction_range + res_dict["empty_prompt"] = empty_prompt + if config["task_instructions"]: + res_dict["final_input_prompt"] = ( + config["task_instructions"].strip() + "\n\n" + empty_prompt + ) + else: + res_dict["final_input_prompt"] = empty_prompt + + res_dict["gt_content"] = options[ord(sample["answer"].upper()) - ord("A")] + else: + empty_prompt_sample_structure = config["short_ans_example_format"] + empty_prompt = empty_prompt_sample_structure.format(question) + res_dict = {} + res_dict["empty_prompt"] = empty_prompt + if config["task_instructions"]: + res_dict["final_input_prompt"] = ( + config["task_instructions"].strip() + "\n\n" + empty_prompt + ) + else: + res_dict["final_input_prompt"] = empty_prompt + res_dict["gt_content"] = sample["answer"] + + res_dict.update(sample) + return res_dict diff --git a/benchmark/mmmu/eval_utils.py b/benchmark/mmmu/eval_utils.py new file mode 100644 index 000000000..ca0e87c6a --- /dev/null +++ b/benchmark/mmmu/eval_utils.py @@ -0,0 +1,649 @@ +"""Response Parsing and Evaluation for various models""" + +import argparse +import dataclasses +import json +import os +import pprint +import random +import re +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Dict, Optional + +import numpy as np +import torch +from data_utils import ( + CAT_SHORT2LONG, + DOMAIN_CAT2SUB_CAT, + construct_prompt, + load_yaml, + process_single_sample, +) +from datasets import concatenate_datasets, load_dataset +from tqdm import tqdm + + +@dataclasses.dataclass +class EvalArgs: + seed: int = 42 + split: str = "validation" + image_pixels_limit: int = -1 + result_filename: str = "" + prompt_format_file: str = "prompt_format.yaml" + dataset_path: str = "MMMU/MMMU" + extra_request_body: Optional[str] = None + profile: bool = False + profile_number: int = 5 + concurrency: int = 1 + response_answer_regex: str = "(.*)" + lora_path: Optional[str] = None + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument( + "--result-filename", + type=str, + default=EvalArgs.result_filename, + help="The filename to save the evaluation results.", + ) + parser.add_argument( + "--image-pixels-limit", + type=int, + default=EvalArgs.image_pixels_limit, + help="The maximum number of pixels allowed for an image. If an image exceeds this limit, it will be skipped during evaluation.", + ) + parser.add_argument( + "--dataset-path", + type=str, + default=EvalArgs.dataset_path, + help="path to the dataset", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--prompt-format-file", + type=str, + help="The path to the prompt format of mmmu. If not, a default format llava_config.yaml will be used", + ) + parser.add_argument( + "--split", + type=str, + default=EvalArgs.split, + help='Split of the dataset to use for evaluation. Default is "validation".', + ) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + default=EvalArgs.extra_request_body, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + parser.add_argument( + "--profile", action="store_true", help="enable mmmu profile" + ) + parser.add_argument( + "--profile-number", + type=int, + default=EvalArgs.profile_number, + help="Number of samples to profile. If not set, will profile all samples.", + ) + parser.add_argument( + "--concurrency", + type=int, + default=EvalArgs.concurrency, + help="Number of concurrent requests to make during evaluation. Default is 1, which means no concurrency.", + ) + parser.add_argument( + "--response-answer-regex", + type=str, + default=EvalArgs.response_answer_regex, + help="Specific regex to capture the answer from the response, string", + ) + parser.add_argument( + "--lora-path", + type=str, + default=EvalArgs.lora_path, + help="Specify the LoRA path to use for evaluation. If specified, the value will be specified in the body of every request as `lora-path`.", + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + attrs = [attr.name for attr in dataclasses.fields(cls)] + return cls(**{attr: getattr(args, attr) for attr in attrs}) + + +def set_seed(seed_value): + """ + Set the seed for PyTorch (both CPU and CUDA), Python, and NumPy for reproducible results. + + :param seed_value: An integer value to be used as the seed. + """ + torch.manual_seed(seed_value) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed_value) + torch.cuda.manual_seed_all(seed_value) # For multi-GPU setups + random.seed(seed_value) + np.random.seed(seed_value) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def prepare_samples(eval_args: EvalArgs): + print("Preparing samples...") + # Build prompts + set_seed(eval_args.seed) + + prompt_format_file = ( + eval_args.prompt_format_file + if eval_args.prompt_format_file is not None + else os.path.join(os.path.dirname(__file__), "prompt_format.yaml") + ) + # load config and process to one value + eval_args.config = load_yaml(prompt_format_file) + for key, value in eval_args.config.items(): + if key != "eval_params" and type(value) == list: + assert len(value) == 1, "key {} has more than one value".format(key) + eval_args.config[key] = value[0] + + # run for each subject in parallel + sub_dataset_list = [] + subjects = list(CAT_SHORT2LONG.values()) # Get a fixed list of subjects + + print(f"Loading datasets for {len(subjects)} subjects...") + with ThreadPoolExecutor() as executor: + # Submit all load_dataset tasks + future_to_subject = { + executor.submit( + load_dataset, eval_args.dataset_path, subject, split=eval_args.split + ): subject + for subject in subjects + } + + # Collect results as they complete + results = {} + for future in tqdm( + as_completed(future_to_subject), + total=len(subjects), + desc="Loading datasets", + ): + subject = future_to_subject[future] + try: + results[subject] = future.result() + except Exception as exc: + print(f"{subject} generated an exception: {exc}") + + # Ensure datasets are added in the original order for consistency + for subject in subjects: + if subject in results: + sub_dataset_list.append(results[subject]) + else: + # Handle cases where a dataset failed to load (optional, depends on desired behavior) + print(f"Warning: Dataset for subject '{subject}' could not be loaded.") + + # merge all dataset + dataset = concatenate_datasets(sub_dataset_list) + + # Prepare images in parallel + images_path = os.path.expanduser("~/.cache/mmmu/images") + os.makedirs(images_path, exist_ok=True) + print(f"Saving images to: {images_path}") + + samples = [] + skip_count = 0 + + def process_sample(i, sample): + sample = process_single_sample(sample) + sample = construct_prompt(sample, eval_args.config) + image = sample["image"] + width, height = image.size + if 0 < eval_args.image_pixels_limit <= width * height: + return None, True + # Use a unique identifier for the image path to avoid potential collisions if indices reset + image_path = f"{images_path}/image_{sample['id']}.png" + if not os.path.exists(image_path): + image.save(image_path) + sample["image_path"] = image_path + return sample, False + + print("Processing samples...") + with ThreadPoolExecutor() as executor: + # Pass the sample itself to process_sample, index is less reliable now + futures = [ + executor.submit( + process_sample, i, sample + ) # Keep index i for tqdm maybe? Or remove it. Let's keep it for now. + for i, sample in enumerate(dataset) + ] + for future in tqdm( + as_completed(futures), total=len(dataset), desc="Processing samples" + ): + sample, skipped = future.result() + if skipped: + skip_count += 1 + elif sample: + samples.append(sample) + + samples.sort(key=lambda x: x["final_input_prompt"]) + + print( + f"Skipping {skip_count} samples with large images, {round((float(skip_count) / len(dataset)) * 100, 2)}% of dataset" + ) + print("Samples have been prepared") + return samples + + +def get_sampling_params(eval_args): + max_new_tokens = 30 + temperature = 0.001 + + extra_request_body = {} + if eval_args.extra_request_body: + extra_request_body = json.loads(eval_args.extra_request_body) + + return { + "temperature": temperature, + "max_new_tokens": max_new_tokens, + **extra_request_body, + } + + +# ----------- Process Multi-choice ------------- +def parse_multi_choice_response(response, all_choices, index2ans): + """ + Parse the prediction from the generated response. + Return the predicted index e.g., A, B, C, D. + """ + for char in [",", ".", "!", "?", ";", ":", "'"]: + response = response.strip(char) + response = " " + response + " " # add space to avoid partial match + + index_ans = True + ans_with_brack = False + candidates = [] + for choice in all_choices: # e.g., (A) (B) (C) (D) + if f"({choice})" in response: + candidates.append(choice) + ans_with_brack = True + + if len(candidates) == 0: + for choice in all_choices: # e.g., A B C D + if f" {choice} " in response: + candidates.append(choice) + + # if all above doesn't get candidates, check if the content is larger than 5 tokens and try to parse the example + if len(candidates) == 0 and len(response.split()) > 5: + for index, ans in index2ans.items(): + if ans.lower() in response.lower(): + candidates.append(index) + index_ans = False # it's content ans. + + if len(candidates) == 0: # still not get answer, randomly choose one. + pred_index = random.choice(all_choices) + elif len(candidates) > 1: + start_indexes = [] + if index_ans: + if ans_with_brack: + for can in candidates: + index = response.rfind(f"({can})") + start_indexes.append(index) # -1 will be ignored anyway + # start_indexes = [generated_response.index(f'({can})') for can in candidates] + else: + for can in candidates: + index = response.rfind(f" {can} ") + start_indexes.append(index) + else: + for can in candidates: + index = response.lower().rfind(index2ans[can].lower()) + start_indexes.append(index) + # get the last one + pred_index = candidates[np.argmax(start_indexes)] + else: # if only one candidate, use it. + pred_index = candidates[0] + + return pred_index + + +# ----------- Process Open ------------- +def check_is_number(string): + """ + Check if the given string a number. + """ + try: + float(string.replace(",", "")) + return True + except ValueError: + # check if there's comma inside + return False + + +def normalize_str(string): + """ + Normalize the str to lower case and make them float numbers if possible. + """ + # check if characters in the string + + # if number, numerize it. + string = string.strip() + + is_number = check_is_number(string) + + if is_number: + string = string.replace(",", "") + string = float(string) + # leave 2 decimal + string = round(string, 2) + return [string] + else: # it's likely to be a string + # lower it + string = string.lower() + if len(string) == 1: + return [" " + string, string + " "] # avoid trivial matches + return [string] + + +def extract_numbers(string): + """ + Exact all forms of numbers from a string with regex. + """ + # Pattern for numbers with commas + pattern_commas = r"-?\b\d{1,3}(?:,\d{3})+\b" + # Pattern for scientific notation + pattern_scientific = r"-?\d+(?:\.\d+)?[eE][+-]?\d+" + # Pattern for simple numbers without commas + pattern_simple = r"-?(?:\d+\.\d+|\.\d+|\d+\b)(?![eE][+-]?\d+)(?![,\d])" + + # Extract numbers with commas + numbers_with_commas = re.findall(pattern_commas, string) + # Extract numbers in scientific notation + numbers_scientific = re.findall(pattern_scientific, string) + # Extract simple numbers without commas + numbers_simple = re.findall(pattern_simple, string) + + # Combine all extracted numbers + all_numbers = numbers_with_commas + numbers_scientific + numbers_simple + return all_numbers + + +def parse_open_response(response): + """ + Parse the prediction from the generated response. + Return a list of predicted strings or numbers. + """ + + # content = content.strip("\n").strip(".").strip(" ") + def get_key_subresponses(response): + key_responses = [] + response = response.strip().strip(".").lower() + sub_responses = re.split(r"\.\s(?=[A-Z])|\n", response) + indicators_of_keys = [ + "could be ", + "so ", + "is ", + "thus ", + "therefore ", + "final ", + "answer ", + "result ", + ] + key_responses = [] + for index, resp in enumerate(sub_responses): + # if last one, accept it's an equation (the entire response can be just one sentence with equation) + if index == len(sub_responses) - 1: + indicators_of_keys.extend(["="]) + shortest_key_response = None # the shortest response that may contain the answer (tail part of the response) + for indicator in indicators_of_keys: + if indicator in resp: + if not shortest_key_response: + shortest_key_response = resp.split(indicator)[-1].strip() + else: + if len(resp.split(indicator)[-1].strip()) < len( + shortest_key_response + ): + shortest_key_response = resp.split(indicator)[-1].strip() + # key_responses.append(resp.split(indicator)[1].strip()) + + if shortest_key_response: + # and it's not trivial + if shortest_key_response.strip() not in [ + ":", + ",", + ".", + "!", + "?", + ";", + ":", + "'", + ]: + key_responses.append(shortest_key_response) + if len(key_responses) == 0: # did not found any + return [response] + return key_responses + + # pdb.set_trace() + key_responses = get_key_subresponses(response) + + pred_list = key_responses.copy() # keep the original string response + for resp in key_responses: + pred_list.extend(extract_numbers(resp)) + + tmp_pred_list = [] + for i in range(len(pred_list)): + tmp_pred_list.extend(normalize_str(pred_list[i])) + pred_list = tmp_pred_list + + # remove duplicates + pred_list = list(set(pred_list)) + + return pred_list + + +# ----------- Evaluation ------------- + + +def eval_multi_choice(gold_i, pred_i): + """ + Evaluate a multiple choice instance. + """ + correct = False + # only they are exactly the same, we consider it as correct + if isinstance(gold_i, list): + for answer in gold_i: + if answer == pred_i: + correct = True + break + else: # gold_i is a string + if gold_i == pred_i: + correct = True + return correct + + +def eval_open(gold_i, pred_i): + """ + Evaluate an open question instance + """ + correct = False + if isinstance(gold_i, list): + # use float to avoid trivial matches + norm_answers = [] + for answer in gold_i: + norm_answers.extend(normalize_str(answer)) + else: + norm_answers = normalize_str(gold_i) + for pred in pred_i: # pred is already normalized in parse response phase + if isinstance(pred, str): # if it's a string, then find if ans in the pred_i + for norm_ans in norm_answers: + # only see if the string answer in the string pred + if isinstance(norm_ans, str) and norm_ans in pred: + if not correct: + correct = True + break + else: # it's a float number + if pred in norm_answers: + if not correct: + correct = True + break + return correct + + +# ----------- Batch Evaluation ------------- +def evaluate(samples): + """ + Batch evaluation for multiple choice and open questions. + """ + pred_correct = 0 + judge_dict = dict() + for sample in samples: + gold_i = sample["answer"] + pred_i = sample["parsed_pred"] + if sample["question_type"] == "multiple-choice": + correct = eval_multi_choice(gold_i, pred_i) + else: # open question + correct = eval_open(gold_i, pred_i) + + if correct: + judge_dict[sample["id"]] = "Correct" + pred_correct += 1 + else: + # print(f"Wrong! expected {pred_i}, answered with {gold_i}") + judge_dict[sample["id"]] = "Wrong" + + if len(samples) == 0: + return {"acc": 0} + return judge_dict, {"acc": pred_correct / len(samples)} + + +# ----------- Calculate Accuracy ------------- +def calculate_ins_level_acc(results: Dict): + """Calculate the instruction level accuracy for given Subject results""" + acc = 0 + ins_num = 0 + for cat_results in results.values(): + acc += cat_results["acc"] * cat_results["num_example"] + ins_num += cat_results["num_example"] + if ins_num == 0: + return 0 + return acc / ins_num + + +def process_result(response, sample, answer_dict, out_samples): + if response is None: + return + if sample["question_type"] == "multiple-choice": + pred_ans = parse_multi_choice_response( + response, sample["all_choices"], sample["index2ans"] + ) + else: # open question + pred_ans = response + + out_samples[sample["id"]] = pred_ans + + # set ground truth answer + answer_dict[sample["id"]] = { + "question_type": sample["question_type"], + "ground_truth": sample["answer"], + } + + +def eval_result(model_answer_path, answer_dict, eval_output_path=None): + if eval_output_path is None: + eval_output_path = model_answer_path + print("Evaluating...") + output_dict = json.load(open(model_answer_path)) + # answer_dict = json.load(open(answer_path)) + + # group by category + output_dict_w_cat = {} + for data_id, parsed_pred in output_dict.items(): + category = "_".join(data_id.split("_")[1:-1]) + if category not in output_dict_w_cat: + output_dict_w_cat.update({category: {}}) + output_dict_w_cat[category].update({data_id: parsed_pred}) + + # group by category + answer_dict_w_cat = {} + for data_id, parsed_pred in answer_dict.items(): + category = "_".join(data_id.split("_")[1:-1]) + if category not in answer_dict_w_cat: + answer_dict_w_cat.update({category: {}}) + answer_dict_w_cat[category].update({data_id: parsed_pred}) + + evaluation_result = {} + + for category in CAT_SHORT2LONG.values(): + # print("Evaluating: {}".format(category)) + # get cat_outputs and cat_answers + try: + cat_outputs = output_dict_w_cat[category] + cat_answers = answer_dict_w_cat[category] + except KeyError: + # print("Skipping {} for not found".format(category)) + continue + + exampels_to_eval = [] + for data_id, parsed_pred in cat_outputs.items(): + question_type = cat_answers[data_id]["question_type"] + if question_type != "multiple-choice": + parsed_pred = parse_open_response( + parsed_pred + ) # mainly for type consistency (make it number, etc.) + else: + parsed_pred = parsed_pred + + exampels_to_eval.append( + { + "id": data_id, + "question_type": question_type, + "answer": cat_answers[data_id]["ground_truth"], + "parsed_pred": parsed_pred, + } + ) + + judge_dict, metric_dict = evaluate(exampels_to_eval) + metric_dict.update({"num_example": len(exampels_to_eval)}) + + evaluation_result[category] = metric_dict + + printable_results = {} + # pdb.set_trace() + # add domain Subject + for domain, in_domain_cats in DOMAIN_CAT2SUB_CAT.items(): + in_domain_cat_results = {} + for cat_name in in_domain_cats: # use the order in DOMAIN_CAT2SUB_CAT + if cat_name in evaluation_result.keys(): + in_domain_cat_results[cat_name] = evaluation_result[cat_name] + else: + pass + in_domain_ins_acc = calculate_ins_level_acc(in_domain_cat_results) + in_domain_data_num = sum( + [ + cat_results["num_example"] + for cat_results in in_domain_cat_results.values() + ] + ) + printable_results["Overall-" + domain] = { + "num": int(in_domain_data_num), + "acc": round(in_domain_ins_acc, 3), + } + # add sub category + for cat_name, cat_results in in_domain_cat_results.items(): + printable_results[cat_name] = { + "num": int(cat_results["num_example"]), + "acc": round(cat_results["acc"], 3), + } + + # table.append(["-----------------------------", "-----", "----"]) + all_ins_acc = calculate_ins_level_acc(evaluation_result) + overall_acc = round(all_ins_acc, 3) + printable_results["Overall"] = { + "num": sum( + [cat_results["num_example"] for cat_results in evaluation_result.values()] + ), + "acc": overall_acc, + } + pprint.pprint(printable_results) + out = eval_output_path + with open(out, "w", encoding="utf-8") as outfile: + json.dump(printable_results, outfile) + print(f"eval out saved to {out}") + + print(f"Overall accuracy: {overall_acc}") diff --git a/benchmark/mmmu/internvl_utils.py b/benchmark/mmmu/internvl_utils.py new file mode 100644 index 000000000..44c62c99a --- /dev/null +++ b/benchmark/mmmu/internvl_utils.py @@ -0,0 +1,94 @@ +# copy from https://huggingface.co/OpenGVLab/InternVL3-1B +import torch +import torchvision.transforms as T +from PIL import Image +from torchvision.transforms.functional import InterpolationMode + +IMAGENET_MEAN = (0.485, 0.456, 0.406) +IMAGENET_STD = (0.229, 0.224, 0.225) + + +def build_transform(input_size): + MEAN, STD = IMAGENET_MEAN, IMAGENET_STD + transform = T.Compose( + [ + T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img), + T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC), + T.ToTensor(), + T.Normalize(mean=MEAN, std=STD), + ] + ) + return transform + + +def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size): + best_ratio_diff = float("inf") + best_ratio = (1, 1) + area = width * height + for ratio in target_ratios: + target_aspect_ratio = ratio[0] / ratio[1] + ratio_diff = abs(aspect_ratio - target_aspect_ratio) + if ratio_diff < best_ratio_diff: + best_ratio_diff = ratio_diff + best_ratio = ratio + elif ratio_diff == best_ratio_diff: + if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]: + best_ratio = ratio + return best_ratio + + +def dynamic_preprocess( + image, min_num=1, max_num=12, image_size=448, use_thumbnail=False +): + orig_width, orig_height = image.size + aspect_ratio = orig_width / orig_height + + # calculate the existing image aspect ratio + target_ratios = set( + (i, j) + for n in range(min_num, max_num + 1) + for i in range(1, n + 1) + for j in range(1, n + 1) + if i * j <= max_num and i * j >= min_num + ) + target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1]) + + # find the closest aspect ratio to the target + target_aspect_ratio = find_closest_aspect_ratio( + aspect_ratio, target_ratios, orig_width, orig_height, image_size + ) + + # calculate the target width and height + target_width = image_size * target_aspect_ratio[0] + target_height = image_size * target_aspect_ratio[1] + blocks = target_aspect_ratio[0] * target_aspect_ratio[1] + + # resize the image + resized_img = image.resize((target_width, target_height)) + processed_images = [] + for i in range(blocks): + box = ( + (i % (target_width // image_size)) * image_size, + (i // (target_width // image_size)) * image_size, + ((i % (target_width // image_size)) + 1) * image_size, + ((i // (target_width // image_size)) + 1) * image_size, + ) + # split the image + split_img = resized_img.crop(box) + processed_images.append(split_img) + assert len(processed_images) == blocks + if use_thumbnail and len(processed_images) != 1: + thumbnail_img = image.resize((image_size, image_size)) + processed_images.append(thumbnail_img) + return processed_images + + +def load_image(image_file, input_size=448, max_num=12): + image = Image.open(image_file).convert("RGB") + transform = build_transform(input_size=input_size) + images = dynamic_preprocess( + image, image_size=input_size, use_thumbnail=True, max_num=max_num + ) + pixel_values = [transform(image) for image in images] + pixel_values = torch.stack(pixel_values) + return pixel_values diff --git a/benchmark/mmmu/prompt_format.yaml b/benchmark/mmmu/prompt_format.yaml new file mode 100644 index 000000000..1a0f7211f --- /dev/null +++ b/benchmark/mmmu/prompt_format.yaml @@ -0,0 +1,15 @@ +task_instructions: +- "" +multi_choice_example_format: +- "{} + +{} + +Answer with the option's letter from the given choices directly." + +short_ans_example_format: +- "{} + +Answer the question using a single word or phrase." +temperature: +- 0 diff --git a/benchmark/mtbench/README.md b/benchmark/mtbench/README.md new file mode 100644 index 000000000..fc37caee9 --- /dev/null +++ b/benchmark/mtbench/README.md @@ -0,0 +1,48 @@ +## Download Dataset + +```sh +wget -O question.jsonl https://raw.githubusercontent.com/lm-sys/FastChat/main/fastchat/llm_judge/data/mt_bench/question.jsonl +``` + +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 80 +``` + +### Benchmark sglang EAGLE +``` +python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-algo EAGLE \ + --speculative-draft-model-path lmsys/sglang-EAGLE-LLaMA3-Instruct-8B --speculative-num-steps 5 \ + --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --dtype float16 --port 30000 +``` + +``` +python3 bench_sglang_eagle.py --num-questions 80 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 80 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 80 --backend lightllm +``` diff --git a/benchmark/mtbench/bench_other.py b/benchmark/mtbench/bench_other.py new file mode 100644 index 000000000..5e579e9a6 --- /dev/null +++ b/benchmark/mtbench/bench_other.py @@ -0,0 +1,111 @@ +import argparse +import json +import os +import time +import uuid +from concurrent.futures import ThreadPoolExecutor + +from fastchat.model import get_conversation_template +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate + + +def load_questions(filename): + questions = [] + with open(filename, "r") as fin: + for line in fin: + obj = json.loads(line) + questions.append(obj) + return questions + + +def write_answers(filename, model_id, questions, answers): + with open(os.path.expanduser(filename), "w") as fout: + for i in range(len(answers)): + ans_json = { + "question_id": questions[i]["question_id"], + "answer_id": uuid.uuid4().hex, + "model_id": model_id, + "choices": { + "index": 0, + "turns": [answers[i][0], answers[i][1]], + }, + "tstamp": time.time(), + } + fout.write(json.dumps(ans_json) + "\n") + + +def main(args): + questions = load_questions(args.question_file) + questions = (questions * 10)[: args.num_questions] + max_tokens = 256 + model_id = "llama-2-chat" + + conv_main = get_conversation_template(model_id) + + # Select backend + call_generate = get_call_generate(args) + + answers = [None] * len(questions) + + def get_answer(i): + conv = conv_main.copy() + cur_answers = [] + for j in range(2): + q = questions[i]["turns"][j] + conv.append_message(conv.roles[0], q) + conv.append_message(conv.roles[1], None) + + prompt = conv.get_prompt() + output = call_generate(prompt, temperature=0, max_tokens=max_tokens).strip() + + cur_answers.append(output) + conv.update_last_message(output) + + answers[i] = cur_answers + + # Run requests + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + latency = time.perf_counter() - tic + + print(f"#questions: {len(questions)}, Latency: {latency:.2f}") + + # Write results + answer_file = args.answer_file or f"tmp_output_{args.backend}.txt" + write_answers(answer_file, model_id, questions, answers) + + with open(args.result_file, "a") as fout: + value = { + "task": "mtbench", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--question-file", type=str, default="question.jsonl") + parser.add_argument("--answer-file", type=str, default=None) + parser.add_argument("--num-questions", type=int, default=80) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/mtbench/bench_sglang.py b/benchmark/mtbench/bench_sglang.py new file mode 100644 index 000000000..0d0545b3a --- /dev/null +++ b/benchmark/mtbench/bench_sglang.py @@ -0,0 +1,99 @@ +import argparse +import json +import os +import time +import uuid + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) + + +def load_questions(filename): + questions = [] + with open(filename, "r") as fin: + for line in fin: + obj = json.loads(line) + questions.append(obj) + return questions + + +def write_answers(filename, model_id, questions, answers): + with open(os.path.expanduser(filename), "w") as fout: + for i in range(len(answers)): + ans_json = { + "question_id": questions[i]["question_id"], + "answer_id": uuid.uuid4().hex, + "model_id": model_id, + "choices": { + "index": 0, + "turns": [answers[i][0], answers[i][1]], + }, + "tstamp": time.time(), + } + fout.write(json.dumps(ans_json) + "\n") + + +@sgl.function +def answer_mt_bench(s, question_1, question_2): + s += sgl.system() + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1")) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2")) + + +def main(args): + # Construct prompts + questions = load_questions(args.question_file)[: args.num_questions] + arguments = [ + {"question_1": q["turns"][0], "question_2": q["turns"][1]} for q in questions + ] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.perf_counter() + rets = answer_mt_bench.run_batch( + arguments, + temperature=0, + max_new_tokens=256, + num_threads=args.parallel, + progress_bar=True, + ) + answers = [[s["answer_1"], s["answer_2"]] for s in rets] + latency = time.perf_counter() - tic + + print(f"#questions: {len(questions)}, Latency: {latency:.2f}") + + # Write results + model_id = backend.model_info["model_path"] + answer_file = args.answer_file or f"tmp_output_{args.backend}.txt" + write_answers(answer_file, model_id, questions, answers) + + with open(args.result_file, "a") as fout: + value = { + "task": "mtbench", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--question-file", type=str, default="question.jsonl") + parser.add_argument("--answer-file", type=str, default=None) + parser.add_argument("--num-questions", type=int, default=80) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/mtbench/bench_sglang_eagle.py b/benchmark/mtbench/bench_sglang_eagle.py new file mode 100644 index 000000000..3eb6036c7 --- /dev/null +++ b/benchmark/mtbench/bench_sglang_eagle.py @@ -0,0 +1,137 @@ +""" +Adapted from https://github.com/chromecast56/sglang/blob/6f145d2eadb93a116134f703358ce76f15381045/benchmark/mtbench/bench_sglang.py + +Benchmark SGLang EAGLE/EAGLE3 Speculative Decoding + +Usage: +python3 benchmark/mtbench/bench_sglang_eagle.py --num-questions 80 --parallel 1 +""" + +import argparse +import json +import os +import time +import uuid + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) + + +def load_questions(filename): + questions = [] + with open(filename, "r") as fin: + for line in fin: + obj = json.loads(line) + questions.append(obj) + return questions + + +def write_answers(filename, model_id, questions, answers): + with open(os.path.expanduser(filename), "w") as fout: + for i in range(len(answers)): + ans_json = { + "question_id": questions[i]["question_id"], + "answer_id": uuid.uuid4().hex, + "model_id": model_id, + "choices": { + "index": 0, + "turns": [answers[i][0], answers[i][1]], + }, + "tstamp": time.time(), + } + fout.write(json.dumps(ans_json) + "\n") + + +@sgl.function +def answer_mt_bench(s, question_1, question_2): + s += sgl.system( + "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information." + ) + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1")) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2")) + + +def main(args): + # Construct prompts + questions = load_questions(args.question_file)[: args.num_questions] + arguments = [ + {"question_1": q["turns"][0], "question_2": q["turns"][1]} for q in questions + ] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.perf_counter() + rets = answer_mt_bench.run_batch( + arguments, + temperature=0, + max_new_tokens=2048, + num_threads=args.parallel, + progress_bar=True, + ) + answers = [[s["answer_1"], s["answer_2"]] for s in rets] + + latency = time.perf_counter() - tic + num_output_tokens = sum( + s.get_meta_info("answer_1")["completion_tokens"] + + s.get_meta_info("answer_2")["completion_tokens"] + for s in rets + ) + + # NOTE: acceptance length is just completion_tokens / spec_verify_ct + # {'id': '3bb9c5ead109488d8ed5ee9cbecaec29', 'finish_reason': {'type': 'length', 'length': 256}, 'prompt_tokens': 37, 'spec_verify_ct': 101, 'completion_tokens': 256, 'cached_tokens': 0} + + output_throughput = num_output_tokens / latency + + has_verify = "spec_verify_ct" in rets[0].get_meta_info("answer_1") + if has_verify: + num_verify_tokens = sum( + s.get_meta_info("answer_1")["spec_verify_ct"] + + s.get_meta_info("answer_2")["spec_verify_ct"] + for s in rets + ) + + accept_length = num_output_tokens / num_verify_tokens + else: + accept_length = 1.0 + + print( + f"#questions: {len(questions)}, Throughput: {output_throughput:.2f} token/s, Acceptance length: {accept_length:.2f}" + ) + + # Write results + model_id = backend.model_info["model_path"] + answer_file = args.answer_file or f"tmp_output_{args.backend}.txt" + write_answers(answer_file, model_id, questions, answers) + + with open(args.result_file, "a") as fout: + value = { + "task": "mtbench", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "throughput": round(output_throughput, 3), + "accept_length": round(accept_length, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--question-file", type=str, default="question.jsonl") + parser.add_argument("--answer-file", type=str, default=None) + parser.add_argument("--num-questions", type=int, default=80) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/multi_chain_reasoning/README.md b/benchmark/multi_chain_reasoning/README.md new file mode 100644 index 000000000..4c9f740f5 --- /dev/null +++ b/benchmark/multi_chain_reasoning/README.md @@ -0,0 +1,49 @@ +## Download data +``` +wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl +``` + +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --schedule-conservativeness 1.3 +``` + +``` +python3 bench_sglang.py --num-questions 64 +python3 bench_sglang.py --num-questions 32 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 64 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 64 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 8 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --num-questions 64 --backend lmql --parallel 1 +``` diff --git a/benchmark/multi_chain_reasoning/bench_other.py b/benchmark/multi_chain_reasoning/bench_other.py new file mode 100644 index 000000000..f361496ad --- /dev/null +++ b/benchmark/multi_chain_reasoning/bench_other.py @@ -0,0 +1,186 @@ +import argparse +import ast +import asyncio +import json +import re +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +prompt_lib = [ + "Let us think step by step.", + "Approach this methodically. Let's dissect the problem into smaller, more manageable parts.", + "It's important to proceed step by step, ensuring accuracy at each stage.", + "Take a deep breath and break this down.", + "A little bit of arithmetic and a logical approach will help us quickly arrive at the solution to this problem.", + "I am extremely good at math.", +] + + +def multi_chain_gsm8k(question, num_chains, call_generate): + s = "Question: " + question + "\n" + # s += call_generate(s + "Answer: " + prompt_lib[0], max_tokens=256, + # stop="Question", temperature=0) + # return s + + comps = [] + for i in range(num_chains): + comps.append( + call_generate( + s + "Answer: " + prompt_lib[i % num_chains], + max_tokens=256, + temperature=0.3, + stop="Question", + ) + ) + + s += "Answer: To answer this question, here are some possible solutions. " + s += "After considering all of them, I will do a majority vote.\n\n" + for i in range(num_chains): + s += f"Solution {i+1}: " + comps[i].strip() + "\n\n" + s += "\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is " + s += call_generate(s, max_tokens=16, temperature=0, stop=None) + return s + + +async def multi_chain_gsm8k_async(question, num_chains, call_generate): + s = "Question: " + question + "\n" + # s += call_generate(s + "Answer: " + prompt_lib[0], max_tokens=256, + # stop="Question", temperature=0) + # return s + + comps = [] + for i in range(num_chains): + comps.append( + await call_generate( + s + "Answer: " + prompt_lib[i % num_chains], + max_tokens=256, + temperature=0.3, + stop="Question", + ) + ) + + s += "Answer: To answer this question, here are some possible solutions. " + s += "After considering all of them, I will do a majority vote.\n\n" + for i in range(num_chains): + s += f"Solution {i+1}: " + comps[i].strip() + "\n\n" + s += "\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is " + s += await call_generate(s, max_tokens=16, temperature=0, stop=None) + return s + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + k = args.num_shot + + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + + states = [None] * len(labels) + + # Select backend + call_generate = get_call_generate(args) + + # Run requests + if args.backend != "lmql": + # Use thread pool + def get_one_answer(i): + answer = multi_chain_gsm8k(questions[i], args.num_chains, call_generate) + states[i] = answer + + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + else: + # Use asyncio + async def get_one_answer_asyncio(i): + answer = await multi_chain_gsm8k_async( + questions[i], args.num_chains, call_generate + ) + states[i] = answer + + tic = time.perf_counter() + loop = asyncio.get_event_loop() + batches = [ + list(range(i, min(i + args.parallel, len(questions)))) + for i in range(0, len(questions), args.parallel) + ] + for bt in tqdm(batches): + tasks = [get_one_answer_asyncio(k) for k in bt] + loop.run_until_complete(asyncio.gather(*tasks)) + + latency = time.perf_counter() - tic + + preds = [] + for i in range(len(states)): + preds.append(get_answer_value(states[i])) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_chain_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shot", type=int, default=0) + parser.add_argument("--num-chains", type=int, default=5) + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=50) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/multi_chain_reasoning/bench_sglang.py b/benchmark/multi_chain_reasoning/bench_sglang.py new file mode 100644 index 000000000..1d3129db2 --- /dev/null +++ b/benchmark/multi_chain_reasoning/bench_sglang.py @@ -0,0 +1,140 @@ +import argparse +import ast +import json +import re +import time + +import numpy as np + +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +prompt_lib = [ + "Let us think step by step.", + "Approach this methodically. Let's dissect the problem into smaller, more manageable parts.", + "It's important to proceed step by step, ensuring accuracy at each stage.", + "Take a deep breath and break this down.", + "A little bit of arithmetic and a logical approach will help us quickly arrive at the solution to this problem.", + "I am extremely good at math.", +] + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + # k = args.num_shot + # few_shot_examples = get_few_shot_examples(lines, k) + + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q} for q in questions] + + num_chains = args.num_chains + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + @sgl.function + def multi_chain_gsm8k(s, question): + s += "Question: " + question + "\n" + # s += "Answer: " + prompt_lib[0] + sgl.gen("answer", max_tokens=256, stop="Question", + # temperature=0) + # return + + forks = s.fork(num_chains) + for i in range(num_chains): + forks[i] += ( + "Answer: " + + prompt_lib[i % num_chains] + + sgl.gen("chain", max_tokens=256, temperature=0.3, stop="Question") + ) + forks.join() + + s += "Answer: To answer this question, here are some possible solutions. " + s += "After considering all of them, I will do a majority vote.\n\n" + for i in range(num_chains): + s += f"Solution {i+1}: " + forks[i]["chain"].strip() + "\n\n" + s += "\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is " + s += sgl.gen("answer", max_tokens=16) + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Select backend + backend = select_sglang_backend(args) + + # Run requests + tic = time.perf_counter() + states = multi_chain_gsm8k.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + preds = [] + for i in range(len(states)): + preds.append(get_answer_value(states[i]["answer"])) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_chain_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shot", type=int, default=0) + parser.add_argument("--num-chains", type=int, default=5) + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=50) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/multi_document_qa/README.md b/benchmark/multi_document_qa/README.md new file mode 100644 index 000000000..09f49c78f --- /dev/null +++ b/benchmark/multi_document_qa/README.md @@ -0,0 +1,47 @@ +## Run benchmark + +### Benchmark sglang +``` +python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-instruct-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 10 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model codellama/CodeLlama-7b-instruct-hf --disable-log-requests --port 21000 --gpu 0.97 +``` + +``` +python3 bench_other.py --backend vllm --num-questions 64 +``` + + +### Benchmark guidance +``` +python3 bench_other.py --backend guidance --num-questions 32 --parallel 1 --n-ctx 11000 --model-path path/to/code-llama/gguf +``` + + + +### Build dataset + +``` +pip install PyPDF2 +python3 build_dataset.py +``` + +```python +import PyPDF2 + +with open('llama2.pdf', 'rb') as file: + reader = PyPDF2.PdfReader(file) + text = '' + for page_num in range(len(reader.pages)): + text += reader.pages[page_num].extract_text() + with open('output.txt', 'w') as text_file: + text_file.write(text) +``` diff --git a/benchmark/multi_document_qa/bench_other.py b/benchmark/multi_document_qa/bench_other.py new file mode 100644 index 000000000..627837c5c --- /dev/null +++ b/benchmark/multi_document_qa/bench_other.py @@ -0,0 +1,114 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +USER_PREFIX = "[INST] " +USER_SUFFIX = " [/INST]" +ASSISTANT_PREFIX = "" +ASSISTANT_SUFFIX = " " + + +def multi_document_qa(docs, question, generate): + s = USER_PREFIX + s += "Please answer a question according to given documents.\n" + s += "Question:" + question + "Documents begin.\n" + + s += "".join(docs) + + s += "\nDocuments end." + s += ( + "\n\nBased on the above documents, please answer this question:\n" + + question + + "\nAnswer in three words or fewer." + ) + s += USER_SUFFIX + s += ASSISTANT_PREFIX + answer = generate(s, max_tokens=16, stop=None) + return answer + + +def main(args): + lines = read_jsonl(args.data_path) + l = lines[0] + arguments = [] + labels = [] + + num_docs = 10 + if args.backend == "guidance": + num_docs = 7 # due to OOM + + for i in range(len(l["questions"][: args.num_questions])): + arguments.append( + { + "docs": l["documents"][:num_docs], + "question": l["questions"][i], + } + ) + labels.append(l["answers"][i]) + states = [None] * len(arguments) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + def get_one_answer(i): + states[i] = multi_document_qa(generate=call_generate, **arguments[i]) + + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(labels))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(labels)))), + total=len(labels), + ) + ) + + latency = time.perf_counter() - tic + + # Compute accuracy + print(states) + correct = 0 + for s, label in zip(states, labels): + answer = s.lower() + if all(x in answer for x in label.lower().split(" ")): + correct += 1 + accuracy = correct / len(labels) + print(f"Accuracy: {accuracy:.3f}") + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_document_qa", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "accuracy": accuracy, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/multi_document_qa/bench_sglang.py b/benchmark/multi_document_qa/bench_sglang.py new file mode 100644 index 000000000..0b4b0dbc6 --- /dev/null +++ b/benchmark/multi_document_qa/bench_sglang.py @@ -0,0 +1,93 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +@sgl.function +def multi_document_qa(s, docs, question): + s += sgl.user_begin() + s += "Please answer a question according to given documents.\n" + s += "Question:" + question + "Documents begin.\n" + + forks = s.fork(len(docs)) + forks += lambda i: docs[i] + forks.join("concate_and_append") + + s += "\nDocuments end." + s += ( + "\n\nBased on the above documents, please answer this question:\n" + + question + + "\nAnswer in three words or fewer." + ) + s += sgl.user_end() + s += sgl.assistant(sgl.gen("answer", max_tokens=16)) + + +def main(args): + lines = read_jsonl(args.data_path) + l = lines[0] + arguments = [] + labels = [] + for i in range(len(l["questions"][: args.num_questions])): + arguments.append( + { + "docs": l["documents"][:10], + "question": l["questions"][i], + } + ) + labels.append(l["answers"][i]) + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.perf_counter() + states = multi_document_qa.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.perf_counter() - tic + + # Compute accuracy + print([s["answer"] for s in states]) + correct = 0 + for s, label in zip(states, labels): + answer = s["answer"].lower() + if all(x in answer for x in label.lower().split(" ")): + correct += 1 + accuracy = correct / len(labels) + print(f"Accuracy: {accuracy:.3f}") + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_document_qa", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "accuracy": accuracy, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/multi_document_qa/build_dataset.py b/benchmark/multi_document_qa/build_dataset.py new file mode 100644 index 000000000..27df9474f --- /dev/null +++ b/benchmark/multi_document_qa/build_dataset.py @@ -0,0 +1,70 @@ +import json + +import transformers + +content = "\n".join( + open("llama2.txt", "r", encoding="utf-8", errors="ignore").readlines() +) +content = content.replace("\n\n", "\n") + +# Count token +name = "meta-llama/Llama-2-7b-chat-hf" +t = transformers.AutoTokenizer.from_pretrained(name) +print(f"num tokens: {len(t.encode(content))}") + +# Segment +SEP = "\n\n" +parts = content.split(SEP) +print(f"num segments: {len(parts)}") + +segment_len = 1100 + +segments = [] +tmp = [] +tmp_len = 0 +for i in range(len(parts)): + tmp.append(parts[i]) + tmp_len += len(t.encode(parts[i])) + + if tmp_len > segment_len: + segments.append(SEP.join(tmp)) + tmp = [] + tmp_len = 0 + +for i, s in enumerate(segments): + print(i, len(t.encode(segments[i]))) + +# Dump +with open("questions.jsonl", "w") as fout: + fout.write( + json.dumps( + { + "documents": segments[:30], + "questions": [ + "What is the name of the fine-tuned LLMs?", + "Which figure shows the helpfulness human evaluation results for Llama 2-Chat?", + "What is the number of parameters in the largest Llama 2 model?", + "What is the batch size of fine-tuning?", + "Where can we find the details of potential data contamination?", + "What is the full name of MPT?", + "What is the power consumption of RSC in Watt?", + "How many tokens of data do they train on?", + "Which model's release is delayed due to a lack of time to sufficiently red team?", + "Which activation function is used in Llama?", + ], + "answers": [ + "Llama 2 Chat", + "1", + "70 B", + "64", + "A 6", + "MosaicML", + "400", + "2 trillion", + "34 B", + "SwiGLU", + ], + } + ) + + "\n" + ) diff --git a/benchmark/multi_turn_chat/README.md b/benchmark/multi_turn_chat/README.md new file mode 100644 index 000000000..0fb5b21fa --- /dev/null +++ b/benchmark/multi_turn_chat/README.md @@ -0,0 +1,66 @@ +### Benchmark sglang + +Run Llama-7B + +``` +python3 -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +Run Mixtral-8x7B +(When there is a CUDA out-of-memory error, try to reduce the `--mem-fraction-static`) + +``` +python3 -m sglang.launch_server --model-path mistralai/Mixtral-8x7B-Instruct-v0.1 --port 30000 --tp-size 8 +``` + +Benchmark(short output) + +``` +python3 bench_sglang.py --tokenizer meta-llama/Llama-2-7b-chat-hf +``` + +Benchmark(long output) + +``` +python3 bench_sglang.py --tokenizer meta-llama/Llama-2-7b-chat-hf --long +``` + +### Benchmark vLLM + +Run Llama-7B + +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +Run Mixtral-8x7B + +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model mistralai/Mixtral-8x7B-Instruct-v0.1 --disable-log-requests --port 21000 --tensor-parallel-size 8 +``` + +Benchmark(short output) + +``` +python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend vllm +``` + +Benchmark(long output) + +``` +python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend vllm --long +``` + +### Benchmark guidance + +Benchmark Llama-7B (short output) + +``` +python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +Benchmark Llama-7B (long output) + +``` +python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf --long +``` diff --git a/benchmark/multi_turn_chat/bench_other.py b/benchmark/multi_turn_chat/bench_other.py new file mode 100644 index 000000000..9189af5be --- /dev/null +++ b/benchmark/multi_turn_chat/bench_other.py @@ -0,0 +1,93 @@ +import json +import time +from argparse import ArgumentParser +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from data_gen import gen_arguments +from tqdm import tqdm +from vllm.transformers_utils.tokenizer import get_tokenizer + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text + + +def multi_turns(generate, qas): + s = "" + for qa in qas: + s += qa["prompt"] + s += generate(s, max_tokens=qa["new_tokens"]) + + return s + + +def main(args): + print(args) + + tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code) + + multi_qas = gen_arguments(args, tokenizer) + + states = [None] * args.num_qa + + call_generate = partial(get_call_generate(args), temperature=0) + + def get_one_answer(i): + states[i] = multi_turns(generate=call_generate, **multi_qas[i]) + + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(multi_qas))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + rets = list( + tqdm( + executor.map(get_one_answer, list(range(len(multi_qas)))), + total=len(multi_qas), + ) + ) + for _ in rets: + pass + + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_turn_chat", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_qa, + "num_turns": args.turns, + "other": { + "parallel": args.parallel, + "output_mode": "long" if args.long else "short", + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--turns", type=int, default=4) + parser.add_argument("--num-qa", type=int, default=20) + parser.add_argument("--min-len-q", type=int, default=256) + parser.add_argument("--max-len-q", type=int, default=512) + parser.add_argument("--min-len-a", type=int, default=4) + parser.add_argument("--max-len-a", type=int, default=8) + parser.add_argument("--tokenizer", type=str, required=True) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--long", action="store_true") + args = add_common_other_args_and_parse(parser) + + if args.long: + args.min_len_a = 256 + args.max_len_a = 512 + args.num_qa = 20 + main(args) diff --git a/benchmark/multi_turn_chat/bench_sglang.py b/benchmark/multi_turn_chat/bench_sglang.py new file mode 100644 index 000000000..1051bf19e --- /dev/null +++ b/benchmark/multi_turn_chat/bench_sglang.py @@ -0,0 +1,79 @@ +import json +import time +from argparse import ArgumentParser + +from data_gen import gen_arguments +from vllm.transformers_utils.tokenizer import get_tokenizer + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +@sgl.function +def multi_turns(s, qas): + for qa in qas: + s += qa["prompt"] + s += sgl.gen(max_tokens=qa["new_tokens"], ignore_eos=True) + + +def main(args): + tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code) + + multi_qas = gen_arguments(args, tokenizer) + + backend = select_sglang_backend(args) + + tic = time.perf_counter() + states = multi_turns.run_batch( + multi_qas, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + print(f"Latency: {latency:.3f}") + + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_turn_chat", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_qa, + "num_turns": args.turns, + "other": { + "parallel": args.parallel, + "output_mode": "long" if args.long else "short", + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--turns", type=int, default=4) + parser.add_argument("--num-qa", type=int, default=20) + parser.add_argument("--min-len-q", type=int, default=256) + parser.add_argument("--max-len-q", type=int, default=512) + parser.add_argument("--min-len-a", type=int, default=4) + parser.add_argument("--max-len-a", type=int, default=8) + parser.add_argument("--tokenizer", type=str, required=True) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--long", action="store_true") + args = add_common_sglang_args_and_parse(parser) + + if args.long: + args.min_len_a = 256 + args.max_len_a = 512 + args.num_qa = 20 + + print(args) + main(args) diff --git a/benchmark/multi_turn_chat/data_gen.py b/benchmark/multi_turn_chat/data_gen.py new file mode 100644 index 000000000..043c07a76 --- /dev/null +++ b/benchmark/multi_turn_chat/data_gen.py @@ -0,0 +1,29 @@ +import random +import string + +random.seed(42) + + +def gen_prompt(tokenizer, token_num): + cha_set = string.ascii_letters + string.digits + ret = "".join(random.choices(cha_set, k=token_num)) + while len(tokenizer(ret).input_ids) < token_num: + ret += random.choice(cha_set) + return ret + + +def gen_arguments(args, tokenizer): + multi_qas = [{"qas": []} for _ in range(args.num_qa)] + for i in range(args.num_qa): + qas = multi_qas[i]["qas"] + for _ in range(args.turns): + prompt_len = random.randint(args.min_len_q, args.max_len_q) + new_tokens = random.randint(args.min_len_a, args.max_len_a) + qas.append( + { + "prompt": gen_prompt(tokenizer, prompt_len), + "new_tokens": new_tokens, + } + ) + + return multi_qas diff --git a/benchmark/multi_turn_chat/long_prompt_multi_turn.py b/benchmark/multi_turn_chat/long_prompt_multi_turn.py new file mode 100644 index 000000000..bda5bb9cc --- /dev/null +++ b/benchmark/multi_turn_chat/long_prompt_multi_turn.py @@ -0,0 +1,129 @@ +import json +import random +import time +from argparse import ArgumentParser +from pathlib import Path + +from tqdm import tqdm + +import sglang as sgl +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +def gen_prompt(tokenizer, token_num): + all_available_tokens = list(tokenizer.get_vocab().values()) + selected_tokens = random.choices(all_available_tokens, k=token_num) + ret = tokenizer.decode(selected_tokens) + return ret + + +def get_cache_path(args): + # Create cache directory under ~/.cache/sglang + cache_dir = Path.home() / ".cache" / "sglang" + + # Create a unique cache filename based on the arguments that affect generation + cache_key = f"qa_{args.num_qa}_{args.turns}_{args.system_prompt_len}_{args.len_q}_{args.len_a}_{args.tokenizer.replace('/', '_')}.json" + return cache_dir / cache_key + + +def gen_arguments(args, tokenizer): + cache_path = get_cache_path(args) + + # Try to load from cache first + if cache_path.exists(): + print(f"Loading cached arguments from {cache_path}") + with open(cache_path, "r") as f: + return json.load(f) + + print("Generating new arguments...") + # First progress bar for system prompts + multi_qas = [] + for _ in tqdm(range(args.num_qa), desc="Generating system prompts"): + multi_qas.append( + {"system_prompt": gen_prompt(tokenizer, args.system_prompt_len), "qas": []} + ) + + # Nested progress bars for QA pairs + for i in tqdm(range(args.num_qa), desc="Generating QA pairs"): + qas = multi_qas[i]["qas"] + for j in range(args.turns): + qas.append( + { + "prompt": gen_prompt(tokenizer, args.len_q), + "new_tokens": args.len_a, + } + ) + + # Save to cache + cache_path.parent.mkdir(parents=True, exist_ok=True) + with open(cache_path, "w") as f: + json.dump(multi_qas, f) + print(f"Cached arguments saved to {cache_path}") + + return multi_qas + + +@sgl.function +def multi_turns(s, system_prompt, qas): + s += system_prompt + + for i, qa in enumerate(qas): + s += qa["prompt"] + s += sgl.gen(max_tokens=qa["new_tokens"], ignore_eos=True) + + +def main(args): + tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code) + + multi_qas = gen_arguments(args, tokenizer) + + backend = select_sglang_backend(args) + + tic = time.perf_counter() + states = multi_turns.run_batch( + multi_qas, + temperature=0, + backend=backend, + num_threads="auto", + progress_bar=True, + ) + latency = time.perf_counter() - tic + + print(f"Latency: {latency:.3f}") + + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_turn_system_prompt_chat", + "backend": args.backend, + "latency": round(latency, 3), + "num_requests": args.num_qa, + "num_turns": args.turns, + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--turns", type=int, default=8) + parser.add_argument("--num-qa", type=int, default=128) + parser.add_argument("--system-prompt-len", type=int, default=2048) + parser.add_argument("--len-q", type=int, default=32) + parser.add_argument("--len-a", type=int, default=128) + parser.add_argument( + "--tokenizer", type=str, default="meta-llama/Meta-Llama-3-8B-Instruct" + ) + parser.add_argument("--trust-remote-code", action="store_true") + args = add_common_sglang_args_and_parse(parser) + + print(args) + main(args) diff --git a/benchmark/prefill_only/bench_embeddings.py b/benchmark/prefill_only/bench_embeddings.py new file mode 100644 index 000000000..ca66c85a3 --- /dev/null +++ b/benchmark/prefill_only/bench_embeddings.py @@ -0,0 +1,148 @@ +""" +SGLang Embeddings Benchmark Script + +This script benchmarks SGLang's /v1/embeddings API performance using HTTP requests. + +Features: +- HTTP-only implementation +- Uses /v1/embeddings API endpoint directly +- Configurable RPS, duration, and batch sizes +- Progress tracking and detailed metrics +- Poisson and constant request distributions + +Usage: +- Update configuration variables at the top of the file +- Ensure SGLang server is running on the configured HTTP_URL +- Run: python bench_embeddings.py +""" + +import asyncio +import logging + +from transformers import AutoTokenizer +from util import ( + BenchmarkConfig, + generate_text_with_token_count, + run_benchmark_main, + run_generic_benchmark, +) + +# Configure logging +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + +############################################################################### +# CONFIG +############################################################################### +# Create benchmark configuration +config = BenchmarkConfig() +config.rps_values = [500] +config.duration_secs_values = [60] +config.num_unique_requests = 100 +config.distribution = "POISSON" +config.profile = False +config.freeze_gc = True # Enable GC freeze functionality +# Profiler output directory - by default uses present working directory (pwd) +# Uncomment and customize the line below to override the default location: +# config.profiler_dir = "/sglang-oss-trace" + +# HTTP Configuration +HTTP_URL = "http://localhost:30000/v1/embeddings" + +# Embeddings API Config +EMBEDDINGS_MODEL_PATH = "/Qwen/Qwen3-Embedding-0.6B" +BATCH_SIZE = [1] # Number of items per request (batch size) + +# Configurable input token length +EMBEDDINGS_INPUT_TOKENS = 500 # Default token length + +# Load tokenizer once for embeddings text generation +print("Loading tokenizer for embeddings input generation...") +embeddings_tokenizer = AutoTokenizer.from_pretrained(EMBEDDINGS_MODEL_PATH) + +# Generate input text with the specified token length using pre-loaded tokenizer +EMBEDDINGS_INPUT_TEXT = generate_text_with_token_count( + EMBEDDINGS_MODEL_PATH, + EMBEDDINGS_INPUT_TOKENS, + config.special_replicated_token, + tokenizer=embeddings_tokenizer, +) + + +############################################################################### +# REQUEST GENERATION (in parallel) +############################################################################### +def build_embeddings_request(index: int, item_count: int) -> tuple: + """Build a single embeddings request.""" + try: + # For embeddings, input can be a string or list of strings + if item_count == 1: + input_data = EMBEDDINGS_INPUT_TEXT + else: + input_data = [EMBEDDINGS_INPUT_TEXT for _ in range(item_count)] + req = { + "input": input_data, + "model": EMBEDDINGS_MODEL_PATH, + } + return (index, req) + except Exception as e: + logger.error(f"Error building request {index}: {e}") + return (index, None) + + +def validate_embeddings_response(response_data: dict) -> bool: + """Validate embeddings API response.""" + return "data" in response_data + + +def build_warmup_embeddings_request() -> dict: + """Build a warmup request for the embeddings API.""" + return { + "input": EMBEDDINGS_INPUT_TEXT, + "model": EMBEDDINGS_MODEL_PATH, + } + + +############################################################################### +# MAIN +############################################################################### +async def run_benchmark(rps, duration_secs, item_count): + """Run a single embeddings benchmark with the given RPS value.""" + return await run_generic_benchmark( + rps=rps, + duration_secs=duration_secs, + item_count=item_count, + config=config, + http_url=HTTP_URL, + build_request_func=build_embeddings_request, + response_validator=validate_embeddings_response, + api_name="EMBEDDINGS", + request_description="embeddings requests", + ) + + +async def main(): + additional_info = { + "Input text length": f"{EMBEDDINGS_INPUT_TOKENS} tokens", + "Input text preview": ( + EMBEDDINGS_INPUT_TEXT[:100] + "..." + if len(EMBEDDINGS_INPUT_TEXT) > 100 + else EMBEDDINGS_INPUT_TEXT + ), + } + + await run_benchmark_main( + config, + run_benchmark, + "EMBEDDINGS", + HTTP_URL, + BATCH_SIZE, + additional_info, + build_warmup_embeddings_request, + ) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/benchmark/prefill_only/bench_score.py b/benchmark/prefill_only/bench_score.py new file mode 100644 index 000000000..117335eae --- /dev/null +++ b/benchmark/prefill_only/bench_score.py @@ -0,0 +1,192 @@ +""" +SGLang Scoring Benchmark Script + +This script benchmarks SGLang's scoring API performance using HTTP requests. + +Current Features: +- HTTP-only implementation (open source compatible) +- Uses /v1/score API endpoint directly +- Single item scoring with batching support +- Configurable RPS, duration, and batch sizes +- Progress tracking and detailed metrics +- Poisson and constant request distributions + +Usage: +- Update configuration variables at the top of the file +- Ensure SGLang server is running on the configured HTTP_URL +- Run: python bench_score.py +- Each request will contain ITEM_COUNT_VALUES items for batch scoring + +""" + +import asyncio + +from transformers import AutoTokenizer +from util import ( + BenchmarkConfig, + generate_text_with_token_count, + run_benchmark_main, + run_generic_benchmark, +) + +############################################################################### +# CONFIG +############################################################################### +# Create benchmark configuration +config = BenchmarkConfig() +config.rps_values = [160] +config.duration_secs_values = [60] +config.num_unique_requests = 100 +config.distribution = "POISSON" +config.profile = False +config.freeze_gc = True # Enable GC freeze functionality +# Profiler output directory - by default uses present working directory (pwd) +# Uncomment and customize the line below to override the default location: +# config.profiler_dir = "/sglang-oss-trace" + +# HTTP Configuration +HTTP_URL = "http://localhost:30000/v1/score" # Use score API directly + +# Score API Config +# ITEM_COUNT_VALUES determines number of items per score request (batch size) +SCORE_QUERY_TOKENS = 120 +SCORE_ITEM_TOKENS = 180 +SCORE_MODEL_PATH = "Qwen/Qwen3-0.6B" +SCORE_LABEL_TOKEN_IDS = [9454, 2753] # Yes/No token IDs +ITEM_COUNT_VALUES = [10] # Number of items per request + +# Special token to replicate for precise token counting +SPECIAL_REPLICATED_TOKEN = "<|im_start|>" + + +############################################################################### +# REQUEST GENERATION (in parallel) +############################################################################### +def create_score_request_builder(): + """Create a score request builder function with shared tokenizer.""" + # Load tokenizer once here to verify special token and get precise counts + print("Loading tokenizer...") + tokenizer = AutoTokenizer.from_pretrained(SCORE_MODEL_PATH) + + # Verify that our special token produces exactly 1 token + special_token_count = len( + tokenizer.encode(config.special_replicated_token, add_special_tokens=False) + ) + print( + f"Special token '{config.special_replicated_token}' produces " + f"{special_token_count} token(s)" + ) + + def generate_text_with_token_count_local(num_toks): + """Generate text with precise token count using replicated token.""" + return generate_text_with_token_count( + SCORE_MODEL_PATH, + num_toks, + config.special_replicated_token, + tokenizer=tokenizer, + ) + + def build_score_request(index: int, item_count: int) -> tuple: + """Build a single score request.""" + try: + # Generate query and items for score API + query = generate_text_with_token_count_local(SCORE_QUERY_TOKENS) + items = [ + generate_text_with_token_count_local(SCORE_ITEM_TOKENS) + for _ in range(item_count) + ] + + # Return as dict for score API format + score_data = { + "query": query, + "items": items, + "label_token_ids": SCORE_LABEL_TOKEN_IDS, + "model": SCORE_MODEL_PATH, + } + return (index, score_data) + + except Exception as e: + print(f"Error building request {index}: {e}") + return (index, None) + + return build_score_request + + +def validate_score_response(response_data: dict) -> bool: + """Validate score API response.""" + return "scores" in response_data or "logprobs" in response_data + + +def build_warmup_score_request() -> dict: + """Build a warmup request for the score API.""" + # Load tokenizer once for warmup generation + tokenizer = AutoTokenizer.from_pretrained(SCORE_MODEL_PATH) + + warmup_query = generate_text_with_token_count( + SCORE_MODEL_PATH, + SCORE_QUERY_TOKENS, + config.special_replicated_token, + tokenizer=tokenizer, + ) + warmup_items = [ + generate_text_with_token_count( + SCORE_MODEL_PATH, + SCORE_ITEM_TOKENS, + config.special_replicated_token, + tokenizer=tokenizer, + ) + for _ in range(3) + ] + + return { + "query": warmup_query, + "items": warmup_items, + "label_token_ids": SCORE_LABEL_TOKEN_IDS, + "model": SCORE_MODEL_PATH, + # Add missing parameters for consistency with the original warmup + "apply_softmax": True, + "item_first": False, + } + + +############################################################################### +# MAIN +############################################################################### +async def run_benchmark(rps, duration_secs, item_count): + """Run a single benchmark with the given RPS value.""" + # Create the request builder function with shared tokenizer + build_request_func = create_score_request_builder() + + return await run_generic_benchmark( + rps=rps, + duration_secs=duration_secs, + item_count=item_count, + config=config, + http_url=HTTP_URL, + build_request_func=build_request_func, + response_validator=validate_score_response, + api_name="SINGLE_ITEM_SCORING", + request_description="score requests", + ) + + +async def main(): + """Main function that runs benchmarks for all RPS values.""" + additional_info = { + "Query tokens per request": SCORE_QUERY_TOKENS, + "Item tokens per item": SCORE_ITEM_TOKENS, + } + + await run_benchmark_main( + config, + run_benchmark, + "SINGLE_ITEM_SCORING", + HTTP_URL, + ITEM_COUNT_VALUES, + additional_info, + build_warmup_score_request, + ) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/benchmark/prefill_only/util.py b/benchmark/prefill_only/util.py new file mode 100644 index 000000000..0dbc39027 --- /dev/null +++ b/benchmark/prefill_only/util.py @@ -0,0 +1,813 @@ +""" +Common utilities for SGLang benchmark scripts. + +This module contains shared code for benchmarking different SGLang APIs +including scoring, embeddings, and other endpoints. +""" + +import asyncio +import concurrent.futures +import json +import os +import random +from statistics import mean +from typing import Any, Callable, Dict, List, Optional, Tuple + +import aiohttp +import numpy as np +from tqdm import tqdm +from transformers import AutoTokenizer + + +class BenchmarkConfig: + """Configuration for benchmark parameters.""" + + def __init__(self): + # Common benchmark settings + self.server_type = "HTTP" + self.rps_values = [70] + self.duration_secs_values = [60] + self.num_unique_requests = 100 + self.distribution = "POISSON" # Options: "CONSTANT", "POISSON" + self.profile = False + + # Garbage Collection Control + self.freeze_gc = True # Enable/disable garbage collection freezing + + # Profiler configuration + self.profiler_dir = ( + os.getcwd() + ) # Default profiler output directory (current working directory) + + # Special token for text generation + self.special_replicated_token = "<|im_start|>" + + +def generate_text_with_token_count( + model_path: str, + num_tokens: int, + special_token: str = "<|im_start|>", + tokenizer: Optional[Any] = None, +) -> str: + """ + Generate text with precise token count using a replicated token. + + Args: + model_path: Path to the model for tokenizer + num_tokens: Target number of tokens + special_token: Token to replicate + tokenizer: Optional pre-loaded tokenizer to avoid repeated loading + + Returns: + Generated text with approximately the target token count + """ + if tokenizer is None: + tokenizer = AutoTokenizer.from_pretrained(model_path) + + # Verify token count + special_token_count = len(tokenizer.encode(special_token, add_special_tokens=False)) + + if special_token_count == 1: + # Simple case: token maps to exactly 1 token + return special_token * num_tokens + else: + print(f"Special token '{special_token}' produces {special_token_count} tokens") + # Handle case where special token produces multiple tokens + repetitions = (num_tokens + special_token_count - 1) // special_token_count + text = special_token * repetitions + + # Verify we got the expected token count + actual_tokens = len(tokenizer.encode(text, add_special_tokens=False)) + if actual_tokens < num_tokens: + print(f"Warning: Generated {actual_tokens} tokens, expected {num_tokens}") + + return text + + +def setup_profiler(config: BenchmarkConfig, benchmark_name: str) -> None: + """ + Set up profiler environment if profiling is enabled. + + Args: + config: Benchmark configuration + benchmark_name: Name of the benchmark (used in directory path) + """ + if config.profile: + # Create benchmark-specific subdirectory + profiler_path = os.path.join( + config.profiler_dir, benchmark_name.lower().replace("_", "-") + ) + os.environ["SGLANG_TORCH_PROFILER_DIR"] = profiler_path + print(f"Profiler enabled. Output directory: {profiler_path}") + else: + print("Profiler disabled") + + +def prepare_all_requests_parallel( + num_requests: int, + item_count: int, + build_request_func: Callable[[int, int], Tuple[int, Any]], + config: BenchmarkConfig, + description: str = "requests", +) -> List[Any]: + """ + Generic function to generate unique requests in parallel, then reuse them. + + Args: + num_requests: Total number of requests needed + item_count: Number of items per request (batch size) + build_request_func: Function that takes (index, item_count) and returns (index, request_data) + config: Benchmark configuration + description: Description for progress bars + + Returns: + List of request data objects + """ + + def build_request_wrapper(index): + """Wrapper to call the provided build_request_func.""" + try: + return build_request_func(index, item_count) + except Exception as e: + print(f"Error building request {index}: {e}") + return (index, None) + + # Generate only the unique requests + unique_requests = [None] * config.num_unique_requests + max_workers = min(8, os.cpu_count() or 1) # Limit to 8 threads max + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [] + for i in tqdm( + range(config.num_unique_requests), + desc=f"Submitting {description} generation tasks", + ): + future = executor.submit(build_request_wrapper, i) + futures.append(future) + + # Collect results as they complete + for f in tqdm( + concurrent.futures.as_completed(futures), + desc=f"Building unique {description}", + total=config.num_unique_requests, + ): + try: + index, req_data = f.result() + if req_data is not None: + unique_requests[index] = req_data + else: + print(f"Failed to build request {index}") + except Exception as e: + print(f"Error processing request result: {e}") + + # Check if we have any valid requests + valid_requests = [req for req in unique_requests if req is not None] + if not valid_requests: + raise RuntimeError("Failed to generate any valid requests") + + print( + f"Successfully generated {len(valid_requests)} out of " + f"{config.num_unique_requests} unique {description}" + ) + + # Create the full request list by cycling through unique requests + print( + f"Reusing {len(valid_requests)} unique {description} to create " + f"{num_requests} total requests..." + ) + all_requests = [] + for i in tqdm(range(num_requests), desc=f"Reusing {description}"): + unique_index = i % len(valid_requests) + all_requests.append(valid_requests[unique_index]) + + print(f"All {description} prepared.\n") + return all_requests + + +async def sleep_with_distribution(distribution: str, rps: float) -> None: + """ + Sleep according to the specified distribution pattern. + + Args: + distribution: "CONSTANT" or "POISSON" + rps: Requests per second rate + """ + if distribution == "CONSTANT": + interval = 1 / rps + await asyncio.sleep(interval) + elif distribution == "POISSON": + # For Poisson process, inter-arrival times follow exponential distribution + interval = random.expovariate(rps) + await asyncio.sleep(interval) + else: + raise ValueError( + f"Unknown distribution: {distribution}. Use 'CONSTANT' or 'POISSON'." + ) + + +def build_http_request_json(request_data: Any) -> str: + """ + Generic function to build HTTP request JSON. + + Args: + request_data: The data to serialize to JSON + + Returns: + JSON string representation of the request data + """ + return json.dumps(request_data) + + +async def make_http_call( + session: aiohttp.ClientSession, + request_data: Any, + request_id: int, + results_queue: asyncio.Queue, + http_url: str, + response_validator: Callable[[Dict[str, Any]], bool], + api_name: str = "API", +) -> None: + """ + Generic HTTP call function for API requests. + + Args: + session: aiohttp client session + request_data: Data to send in the request + request_id: Unique identifier for this request + results_queue: Queue to put results + http_url: URL to send the request to + response_validator: Function to validate the response JSON + api_name: Name of the API for error messages + """ + try: + start_time = asyncio.get_event_loop().time() + + request_json = build_http_request_json(request_data) + headers = {"Content-Type": "application/json"} + + async with session.post(http_url, data=request_json, headers=headers) as resp: + resp_text = await resp.text() + + if resp.status != 200: + print( + f"[HTTP] {api_name} Request {request_id} failed with status " + f"{resp.status}: {resp_text}" + ) + completion_time = asyncio.get_event_loop().time() + await results_queue.put((request_id, 0, False, completion_time)) + return + + # Parse and validate response + try: + response_data = json.loads(resp_text) + success = response_validator(response_data) + if not success: + print( + f"[HTTP] {api_name} Request {request_id} failed response validation" + ) + except json.JSONDecodeError: + print( + f"[HTTP] {api_name} Request {request_id} failed to parse JSON response" + ) + success = False + + completion_time = asyncio.get_event_loop().time() + elapsed_time = (completion_time - start_time) * 1000 + await results_queue.put((request_id, elapsed_time, success, completion_time)) + + except Exception as e: + print(f"[HTTP] {api_name} Error for request {request_id}: {e}") + completion_time = asyncio.get_event_loop().time() + await results_queue.put((request_id, 0, False, completion_time)) + + +async def send_profile_request( + profile_text: str, http_url: str, session: Optional[aiohttp.ClientSession] = None +) -> None: + """ + Send a profile request (START_PROFILE or STOP_PROFILE) and wait for completion. + + Args: + profile_text: "START_PROFILE" or "STOP_PROFILE" + http_url: Base HTTP URL (will derive profile endpoints from this) + session: Optional aiohttp session to use + """ + try: + if session: + print(f"Sending {profile_text} request via HTTP...") + + # Determine the correct endpoint + if "/v1/" in http_url: + base_url = http_url.rsplit("/v1/", 1)[0] # Remove /v1/xxx + else: + base_url = http_url.rsplit("/", 1)[0] # Remove last path component + + if profile_text == "START_PROFILE": + endpoint_url = f"{base_url}/start_profile" + elif profile_text == "STOP_PROFILE": + endpoint_url = f"{base_url}/stop_profile" + else: + print(f"Unknown profile request: {profile_text}") + return + + headers = {"Content-Type": "application/json"} + + async with session.post(endpoint_url, headers=headers) as resp: + resp_text = await resp.text() + if resp.status == 200: + print(f"{profile_text} request completed") + else: + print( + f"{profile_text} request failed with status " + f"{resp.status}: {resp_text}" + ) + else: + print(f"Cannot send {profile_text} request - missing session") + + except Exception as e: + print(f"Error sending {profile_text} request: {e}") + + +async def call_freeze_gc_http(session: aiohttp.ClientSession, http_url: str) -> None: + """ + Call the /freeze_gc HTTP endpoint. + + Args: + session: aiohttp client session + http_url: Base HTTP URL to derive the freeze_gc endpoint from + """ + try: + # Derive freeze_gc endpoint from the API URL + if "/v1/" in http_url: + freeze_gc_url = http_url.rsplit("/v1/", 1)[0] + "/freeze_gc" + else: + freeze_gc_url = http_url.rsplit("/", 1)[0] + "/freeze_gc" + + print(f"Calling freeze_gc endpoint: {freeze_gc_url}") + + async with session.post(freeze_gc_url) as resp: + if resp.status == 200: + print("freeze_gc called successfully") + else: + resp_text = await resp.text() + print(f"freeze_gc failed with status {resp.status}: {resp_text}") + + except Exception as e: + print(f"Failed to call freeze_gc: {e}") + + +async def send_warmup_requests( + session: aiohttp.ClientSession, + http_url: str, + build_warmup_request_func: Callable[[], Any], + num_warmup: int = 3, +) -> None: + """ + Send warmup requests to HTTP server. + + Args: + session: aiohttp client session + http_url: URL to send warmup requests to + build_warmup_request_func: Function that returns a warmup request object + num_warmup: Number of warmup requests to send + """ + print(f"Sending {num_warmup} HTTP warmup requests...") + + for i in range(num_warmup): + try: + warmup_data = build_warmup_request_func() + request_json = build_http_request_json(warmup_data) + headers = {"Content-Type": "application/json"} + + async with session.post( + http_url, data=request_json, headers=headers + ) as resp: + if resp.status == 200: + print(f"Warmup request {i+1}/{num_warmup} completed successfully") + else: + print( + f"Warmup request {i+1}/{num_warmup} failed with status {resp.status}" + ) + + except Exception as e: + print(f"Warmup request {i+1}/{num_warmup} failed with error: {e}") + + print("HTTP warmup requests completed") + + +async def perform_global_warmup_and_freeze( + config: BenchmarkConfig, + http_url: str, + build_warmup_request_func: Callable[[], Any], +) -> None: + """ + Perform warmup and optionally GC freeze operations once before all benchmark runs. + + Args: + config: Benchmark configuration + http_url: URL for API requests + build_warmup_request_func: Function that returns a warmup request object + """ + print("=" * 80) + print(f"PERFORMING GLOBAL WARMUP{' AND GC FREEZE' if config.freeze_gc else ''}") + print("=" * 80) + + print(f"Performing HTTP warmup{' and GC freeze' if config.freeze_gc else ''}...") + async with aiohttp.ClientSession() as session: + await send_warmup_requests(session, http_url, build_warmup_request_func) + if config.freeze_gc: + await call_freeze_gc_http(session, http_url) + print( + f"HTTP warmup{' and GC freeze' if config.freeze_gc else ''} completed successfully." + ) + + print( + f"Global warmup{' and GC freeze' if config.freeze_gc else ''} operations completed." + ) + print("=" * 80) + + +async def process_results( + results_queue: asyncio.Queue, + num_requests: int, + send_duration: float, + total_duration: float, + rps: int, + duration_secs: int, + item_count: int, + test_start_time: float, + config: BenchmarkConfig, + http_mode: str = "UNKNOWN", +) -> List[Dict[str, Any]]: + """ + Process benchmark results and group them by minute intervals. + + Args: + results_queue: Queue containing result tuples + num_requests: Total number of requests sent + send_duration: Time taken to send all requests + total_duration: Total time for all requests to complete + rps: Target requests per second + duration_secs: Test duration in seconds + item_count: Number of items per request + test_start_time: Start time of the test + config: Benchmark configuration + http_mode: Description of the HTTP mode/API being tested + + Returns: + List of dictionaries containing minute-by-minute results + """ + all_results = [] + + # Collect all results + for _ in range(num_requests): + result = await results_queue.get() + request_id, elapsed_time, success, completion_time = result + all_results.append( + { + "request_id": request_id, + "elapsed_time": elapsed_time, + "success": success, + "completion_time": completion_time, + } + ) + + # Group results by minute intervals + minute_results = [] + num_minutes = int(duration_secs // 60) + (1 if duration_secs % 60 > 0 else 0) + + for minute in range(num_minutes): + minute_start = test_start_time + (minute * 60) + minute_end = test_start_time + ((minute + 1) * 60) + + # Filter results that completed in this minute + minute_data = [ + r for r in all_results if minute_start <= r["completion_time"] < minute_end + ] + + response_times = [r["elapsed_time"] for r in minute_data if r["success"]] + successful_requests = len([r for r in minute_data if r["success"]]) + failed_requests = len([r for r in minute_data if not r["success"]]) + + avg_response_time = mean(response_times) if response_times else 0 + + # Calculate percentiles using numpy + if response_times: + p50 = np.percentile(response_times, 50) + p90 = np.percentile(response_times, 90) + p99 = np.percentile(response_times, 99) + else: + p50 = p90 = p99 = 0 + + minute_result = { + "test_duration_secs": duration_secs, + "minute_interval": minute + 1, + "target_rps": rps, + "item_count": item_count, + "server_type": config.server_type, + "distribution": config.distribution, + "unique_requests": config.num_unique_requests, + "total_requests": len(minute_data), + "successful_requests": successful_requests, + "failed_requests": failed_requests, + "send_duration_secs": send_duration, + "total_duration_secs": total_duration, + "avg_response_time_ms": avg_response_time, + "p50_response_time_ms": p50, + "p90_response_time_ms": p90, + "p99_response_time_ms": p99, + } + + minute_results.append(minute_result) + + print( + f"\nMinute {minute + 1} Summary for RPS {rps}, " + f"Duration {duration_secs}s, Item Count {item_count}:" + ) + print(f" Requests completed in minute: {len(minute_data)}") + print(f" Successful requests: {successful_requests}") + print(f" Failed requests: {failed_requests}") + print(f" Average response time: {avg_response_time:.2f} ms") + print(f" P50 response time: {p50:.2f} ms") + print(f" P90 response time: {p90:.2f} ms") + print(f" P99 response time: {p99:.2f} ms") + + # Print overall summary + all_response_times = [r["elapsed_time"] for r in all_results if r["success"]] + total_successful = len([r for r in all_results if r["success"]]) + total_failed = len([r for r in all_results if not r["success"]]) + + overall_avg = mean(all_response_times) if all_response_times else 0 + if all_response_times: + overall_p50 = np.percentile(all_response_times, 50) + overall_p90 = np.percentile(all_response_times, 90) + overall_p99 = np.percentile(all_response_times, 99) + else: + overall_p50 = overall_p90 = overall_p99 = 0 + + print( + f"\nOverall Summary for RPS {rps}, Duration {duration_secs}s, " + f"Item Count {item_count}:" + ) + print(f" Test duration: {duration_secs} seconds") + print(f" Server type: {config.server_type}") + print(f" HTTP mode: {http_mode}") + print(f" Target RPS: {rps}") + print(f" Item count: {item_count}") + print(f" Distribution: {config.distribution}") + print(f" Unique requests generated: {config.num_unique_requests}") + print(f" Total requests sent: {num_requests}") + print(f" Successful requests: {total_successful}") + print(f" Failed requests: {total_failed}") + print(f" Time to send all requests: {send_duration:.2f} seconds") + print(f" Time for all requests to complete: {total_duration:.2f} seconds") + print(f" Average response time: {overall_avg:.2f} ms") + print(f" P50 response time: {overall_p50:.2f} ms") + print(f" P90 response time: {overall_p90:.2f} ms") + print(f" P99 response time: {overall_p99:.2f} ms\n") + + return minute_results + + +def print_csv_results(all_results: List[Dict[str, Any]]) -> None: + """ + Print benchmark results in CSV format. + + Args: + all_results: List of result dictionaries from process_results + """ + print("\n" + "=" * 80) + print("FINAL CSV RESULTS:") + print("=" * 80) + + # CSV Header + headers = [ + "test_duration_secs", + "minute_interval", + "target_rps", + "item_count", + "server_type", + "distribution", + "unique_requests", + "total_requests", + "successful_requests", + "failed_requests", + "send_duration_secs", + "total_duration_secs", + "avg_response_time_ms", + "p50_response_time_ms", + "p90_response_time_ms", + "p99_response_time_ms", + ] + print(",".join(headers)) + + # CSV Data + for result in all_results: + row = [ + result["test_duration_secs"], + result["minute_interval"], + result["target_rps"], + result["item_count"], + result["server_type"], + result["distribution"], + result["unique_requests"], + result["total_requests"], + result["successful_requests"], + result["failed_requests"], + f"{result['send_duration_secs']:.2f}", + f"{result['total_duration_secs']:.2f}", + f"{result['avg_response_time_ms']:.2f}", + f"{result['p50_response_time_ms']:.2f}", + f"{result['p90_response_time_ms']:.2f}", + f"{result['p99_response_time_ms']:.2f}", + ] + print(",".join(map(str, row))) + + +async def run_benchmark_main( + config: BenchmarkConfig, + run_single_benchmark_func, + benchmark_name: str, + http_url: str, + item_count_values: List[int], + additional_info: Optional[Dict[str, Any]] = None, + build_warmup_request_func: Optional[Callable[[], Any]] = None, +) -> None: + """ + Main benchmark orchestration function. + + Args: + config: Benchmark configuration + run_single_benchmark_func: Async function to run a single benchmark + benchmark_name: Name of the benchmark (e.g., "SCORING", "EMBEDDINGS") + http_url: URL of the API endpoint + item_count_values: List of item counts to test + additional_info: Additional information to print in the header + build_warmup_request_func: Optional function to build warmup requests + """ + total_combinations = ( + len(config.duration_secs_values) + * len(config.rps_values) + * len(item_count_values) + ) + + print( + f"Running benchmarks for {len(config.duration_secs_values)} duration " + f"values, {len(config.rps_values)} RPS values, and " + f"{len(item_count_values)} item count values = " + f"{total_combinations} total combinations" + ) + print(f"Server Type: {config.server_type}") + print(f"HTTP Mode: {benchmark_name}") + print(f"API URL: {http_url}") + + if additional_info: + for key, value in additional_info.items(): + print(f"{key}: {value}") + + print(f"Items per request (batch size): {item_count_values}") + print(f"Profiling Enabled: {config.profile}") + print(f"Duration values: {config.duration_secs_values}") + print(f"RPS values: {config.rps_values}") + print(f"Item count values: {item_count_values}") + print("=" * 80) + + # Set up profiler environment + setup_profiler(config, benchmark_name) + + # Perform global warmup and GC freeze operations if warmup function is provided + if build_warmup_request_func is not None: + await perform_global_warmup_and_freeze( + config, http_url, build_warmup_request_func + ) + + all_results = [] + + for duration_secs in config.duration_secs_values: + for rps in config.rps_values: + for item_count in item_count_values: + result = await run_single_benchmark_func(rps, duration_secs, item_count) + all_results.extend(result) # Extend with minute results + + print_csv_results(all_results) + + +async def run_generic_benchmark( + rps: int, + duration_secs: int, + item_count: int, + config: BenchmarkConfig, + http_url: str, + build_request_func: Callable[[int, int], Tuple[int, Any]], + response_validator: Callable[[Dict[str, Any]], bool], + api_name: str, + request_description: str = "requests", +) -> List[Dict[str, Any]]: + """ + Generic benchmark runner that can be used for different APIs. + + Args: + rps: Requests per second + duration_secs: Duration of the test in seconds + item_count: Number of items per request (batch size) + config: Benchmark configuration + http_url: URL of the API endpoint + build_request_func: Function to build individual requests + response_validator: Function to validate API responses + api_name: Name of the API for logging + request_description: Description for progress bars + + Returns: + List of dictionaries containing minute-by-minute results + """ + num_requests = int(rps * duration_secs) + print( + f"Starting benchmark with RPS={rps}, Duration={duration_secs}s, " + f"Item Count={item_count}, num_requests={num_requests}" + ) + print(f"Server Type: {config.server_type}") + print(f"HTTP Mode: {api_name}") + print(f"Profiling Enabled: {config.profile}") + + # Build requests in parallel (unmeasured) + all_requests = prepare_all_requests_parallel( + num_requests, item_count, build_request_func, config, request_description + ) + + results_queue = asyncio.Queue() + tasks = [] + + # Track timing for sending requests + send_start_time = asyncio.get_event_loop().time() + + # HTTP implementation + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=300) + ) as session: + + # Send START_PROFILE if profiling is enabled + if config.profile: + await send_profile_request("START_PROFILE", http_url, session=session) + + # Add progress bar for sending requests + with tqdm( + total=len(all_requests), + desc=f"Sending HTTP {request_description} at {rps} RPS", + unit="req", + ) as pbar: + for i, request_data in enumerate(all_requests): + request_id = i + 1 + tasks.append( + asyncio.create_task( + make_http_call( + session, + request_data, + request_id, + results_queue, + http_url, + response_validator, + api_name, + ) + ) + ) + + # Update progress bar + pbar.update(1) + + # Throttle based on distribution + if i < len(all_requests) - 1: + await sleep_with_distribution(config.distribution, rps) + + send_end_time = asyncio.get_event_loop().time() + send_duration = send_end_time - send_start_time + + # Wait for all requests to complete with progress tracking + print(f"Waiting for {len(tasks)} HTTP {request_description} to complete...") + with tqdm( + total=len(tasks), desc=f"Completing HTTP {request_description}", unit="req" + ) as completion_pbar: + completed_tasks = [] + for task in asyncio.as_completed(tasks): + await task + completed_tasks.append(task) + completion_pbar.update(1) + + # Send STOP_PROFILE if profiling is enabled + if config.profile: + await send_profile_request("STOP_PROFILE", http_url, session=session) + + completion_end_time = asyncio.get_event_loop().time() + total_duration = completion_end_time - send_start_time + + return await process_results( + results_queue, + num_requests, + send_duration, + total_duration, + rps, + duration_secs, + item_count, + send_start_time, + config, + api_name, + ) diff --git a/benchmark/react/README.md b/benchmark/react/README.md new file mode 100644 index 000000000..51bcaa44a --- /dev/null +++ b/benchmark/react/README.md @@ -0,0 +1,34 @@ +## Run benchmark + +NOTE: This is an implementation for replaying a given trace for throughput/latency benchmark purposes. It is not an actual ReAct agent implementation. + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 100 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 100 --backend vllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 100 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --num-questions 100 --backend lmql --parallel 1 +``` diff --git a/benchmark/react/bench_other.py b/benchmark/react/bench_other.py new file mode 100644 index 000000000..08666662b --- /dev/null +++ b/benchmark/react/bench_other.py @@ -0,0 +1,202 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + + +def get_prompt(question): + prompt = ( + """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: +(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search. +(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage. +(3) Finish[answer], which returns the answer and finishes the task. +Here are some examples. +Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into? +Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area. +Action 1: Search[Colorado orogeny] +Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas. +Thought 2: It does not mention the eastern sector. So I need to look up eastern sector. +Action 2: Lookup[eastern sector] +Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny. +Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range. +Action 3: Search[High Plains] +Observation 3: High Plains refers to one of two distinct land regions: +Thought 4: I need to instead search High Plains (United States). +Action 4: Search[High Plains (United States)] +Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3] +Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft. +Action 5: Finish[1,800 to 7,000 ft] +Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who? +Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after. +Action 1: Search[Milhouse] +Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening. +Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after". +Action 2: Lookup[named after] +Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous. +Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon. +Action 3: Finish[Richard Nixon] +Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture? +Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups. +Action 1: Search[Adam Clayton Powell] +Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito']. +Thought 2: To find the documentary, I can search Adam Clayton Powell (film). +Action 2: Search[Adam Clayton Powell (film)] +Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg. +The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience. +Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture. +Action 3: Finish[The Saimaa Gesture] +Question: What profession does Nicholas Ray and Elia Kazan have in common? +Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common. +Action 1: Search[Nicholas Ray] +Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause. +Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions. +Action 2: Search[Elia Kazan] +Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor. +Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor. +Action 3: Finish[director, screenwriter, actor] +Question: Which magazine was started first Arthur's Magazine or First for Women? +Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first. +Action 1: Search[Arthur's Magazine] +Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century. +Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next. +Action 2: Search[First for Women] +Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989. +Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first. +Action 3: Finish[Arthur's Magazine] +Question: Were Pavel Urysohn and Leonid Levin known for the same type of work? +Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same. +Action 1: Search[Pavel Urysohn] +Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory. +Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work. +Action 2: Search[Leonid Levin] +Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist. +Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work. +Action 3: Finish[yes] +""" + + question + ) + return prompt + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + arguments = [{"question": k, "triplets": v} for l in lines for k, v in l.items()] + + states = [] + + # Select backend + call_generate = get_call_generate(args) + + def run_single_agent(argument): + question = argument["question"] + triplets = argument["triplets"] + prompt = get_prompt(question) + for i in range(1, len(triplets) + 2): + prompt += "Thought " + str(i) + ":" + states.append(prompt) + answer = call_generate( + prompt, max_tokens=200, temperature=0, stop="Observation" + ) + if i > len(triplets): + break + prompt += ( + triplets[i - 1]["thought"] + + "\nAction " + + str(i) + + ":" + + triplets[i - 1]["action"] + + "\nObservation " + + str(i) + + ":" + + triplets[i - 1]["observation"] + + "\n" + ) + + states.append(answer) + + async def run_single_agent_async(argument): + question = argument["question"] + triplets = argument["triplets"] + prompt = get_prompt(question) + for i in range(1, len(triplets) + 2): + prompt += "Thought " + str(i) + ":" + states.append(prompt) + answer = await call_generate( + prompt, max_tokens=200, temperature=0, stop="Observation", max_len=4096 + ) + if i > len(triplets): + break + prompt += ( + triplets[i - 1]["thought"] + + "\nAction " + + str(i) + + ":" + + triplets[i - 1]["action"] + + "\nObservation " + + str(i) + + ":" + + triplets[i - 1]["observation"] + + "\n" + ) + + states.append(answer) + + tic = time.perf_counter() + + if args.backend != "lmql": + if args.parallel == 1: + for arg in tqdm(arguments): + run_single_agent(arg) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(run_single_agent, arguments), total=len(arguments) + ) + ) + + else: + import asyncio + + loop = asyncio.get_event_loop() + batches = [ + [] for _ in range((len(arguments) + args.parallel - 1) // args.parallel) + ] + for i, arg in enumerate(arguments): + batches[i // args.parallel].append(arg) + for bt in tqdm(batches): + tasks = [run_single_agent_async(arg) for arg in bt] + loop.run_until_complete(asyncio.gather(*tasks)) + + latency = time.perf_counter() - tic + + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "ReAct Agents", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": len(arguments), + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="hotpotqa_100.jsonl") + parser.add_argument("--num-questions", type=int, default=10) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/react/bench_sglang.py b/benchmark/react/bench_sglang.py new file mode 100644 index 000000000..331638e9f --- /dev/null +++ b/benchmark/react/bench_sglang.py @@ -0,0 +1,153 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +@sgl.function +def webthink(s, question, triplets): + s += ( + """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: +(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search. +(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage. +(3) Finish[answer], which returns the answer and finishes the task. +Here are some examples. +Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into? +Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area. +Action 1: Search[Colorado orogeny] +Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas. +Thought 2: It does not mention the eastern sector. So I need to look up eastern sector. +Action 2: Lookup[eastern sector] +Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny. +Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range. +Action 3: Search[High Plains] +Observation 3: High Plains refers to one of two distinct land regions: +Thought 4: I need to instead search High Plains (United States). +Action 4: Search[High Plains (United States)] +Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3] +Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft. +Action 5: Finish[1,800 to 7,000 ft] +Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who? +Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after. +Action 1: Search[Milhouse] +Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening. +Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after". +Action 2: Lookup[named after] +Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous. +Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon. +Action 3: Finish[Richard Nixon] +Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture? +Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups. +Action 1: Search[Adam Clayton Powell] +Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito']. +Thought 2: To find the documentary, I can search Adam Clayton Powell (film). +Action 2: Search[Adam Clayton Powell (film)] +Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg. +The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience. +Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture. +Action 3: Finish[The Saimaa Gesture] +Question: What profession does Nicholas Ray and Elia Kazan have in common? +Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common. +Action 1: Search[Nicholas Ray] +Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause. +Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions. +Action 2: Search[Elia Kazan] +Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor. +Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor. +Action 3: Finish[director, screenwriter, actor] +Question: Which magazine was started first Arthur's Magazine or First for Women? +Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first. +Action 1: Search[Arthur's Magazine] +Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century. +Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next. +Action 2: Search[First for Women] +Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989. +Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first. +Action 3: Finish[Arthur's Magazine] +Question: Were Pavel Urysohn and Leonid Levin known for the same type of work? +Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same. +Action 1: Search[Pavel Urysohn] +Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory. +Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work. +Action 2: Search[Leonid Levin] +Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist. +Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work. +Action 3: Finish[yes] +""" + + question + ) + for i in range(1, len(triplets) + 2): + s += "Thought " + str(i) + ":" + # NOTE: This is an implementation for replaying a given trace for benchmark purposes. It is not an actual ReAct agent implementation. + ss = s.fork(1) + ss[0] += sgl.gen(name="thought_action", max_tokens=200, stop="Observation") + ss.join() + # to verify the correctness of output, this should be collected + # print(ss[0]["thought_action"]) + if i > len(triplets): + break + s += ( + triplets[i - 1]["thought"] + + "\nAction " + + str(i) + + ":" + + triplets[i - 1]["action"] + + "\nObservation " + + str(i) + + ":" + + triplets[i - 1]["observation"] + + "\n" + ) + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + arguments = [{"question": k, "triplets": v} for l in lines for k, v in l.items()] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + states = [] + tic = time.perf_counter() + states = webthink.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "ReAct Agents", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": len(arguments), + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="hotpotqa_100.jsonl") + parser.add_argument("--num-questions", type=int, default=10) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/reasoning_benchmark/README.md b/benchmark/reasoning_benchmark/README.md new file mode 100644 index 000000000..f9d26e0c7 --- /dev/null +++ b/benchmark/reasoning_benchmark/README.md @@ -0,0 +1,77 @@ +# Run benchmark + +This benchmark is primarily intended to be used with reasoning models like `DeepSeek-R1` and its distilled models like `DeepSeek-R1-Distill-Qwen-1.5B`. Please use + +```bash +pip install antlr4-python3-runtime +``` + +for `parse_latex` which we use for symbolic equality check. + +## Benchmark sglang + +1. Launch the Server +```bash +python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --port 30000 +``` + +Note that depending on the GPU this benchmark will take quiet some time. To employ data parallelism please use: + +```bash +python3 -m sglang_router.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --port 30000 --dp-size 4 +``` + + +2. Benchmarking + +We use [suggested](https://github.com/deepseek-ai/DeepSeek-R1) parameters of `temperature=0.6`, `top_p=.95`, `max_new_tokens=32768`. The command line argument `num-tries` can be used to evaluate the model multiple times on the same question. We use the suggested `64` from the repo for AIME 2024. For LIMO, we use `8` as the number of tries due to the size of the dataset. + +By default evaluate on LIMO dataset. + +```bash +python3 bench_sglang.py --parallel 256 --num-tries 64 --port 30000 +``` + +Evaluate on AIME 2024 dataset. + +```bash +python3 bench_sglang.py --parallel 256 --port 30000 --data-path Maxwell-Jia/AIME_2024 --question-key Problem --answer-key Answer --num-tries 64 +``` + +Evaluate on [AIME 2025 I dataset](https://huggingface.co/datasets/opencompass/AIME2025). For benchmark result see [here](https://matharena.ai/). + +```bash +python3 bench_sglang.py --parallel 256 --port 30000 --data-path opencompass/AIME2025 --question-key question --answer-key answer --num-tries 64 +``` +## Results + +### Evaluation Results +| Dataset | Num Tries | Accuracy | Reference | Standard Error | +|------------|-----------|----------|-----------|-----------| +| LIMO | 8 | 47.7% | ? | ? | +| AIME 2024 | 64 | 33.2% | 28.9% | 3.4% | +| AIME 2025 I| 64 | 29.9% | 25.0% | ? | + +### Statistic Analysis Results +Set up SGLang engine for statistic analysis, for high efficiency we use `--dp-size 8` for data parallelism: +```bash +python3 -m sglang_router.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --port 30000 --dp-size 8 +``` +**Experiment 1**: +We fixed the number of attempts (num_tries) and conducted multiple runs to assess the consistency of the model's performance. The results show that all recorded accuracies lie within ± one standard error deviation from the mean. This suggests that **our metric serves as an effective upper bound for the deviation of reported accuracy**. + +To collect the accuracy, run the following command 30 times: +```bash +python3 bench_sglang.py --parallel 64 --port 30000 --data-path Maxwell-Jia/AIME_2024 --question-key Problem --answer-key Answer --num-tries 64 +``` + +![acc_hist](figure/Acc_histplot.png) + + +**Experiment 2**: We explored the relationship between the number of attempts (num_tries) and the standard error (SE) by varying num_tries across a range (e.g., 8, 16, 32, ..., 256) and performing a single run for each value. The results demonstrate that as the number of attempts increases, the standard error decreases, leading to **greater stability in answer accuracy**. + +To reveal the relationship, run the command 6 times and adjust the parameter `--num-tries` for each run: +```bash +python3 bench_sglang.py --parallel 64 --port 30000 --data-path Maxwell-Jia/AIME_2024 --question-key Problem --answer-key Answer --num-tries +``` +![SE_num_tries](figure/SE_numtries.png) diff --git a/benchmark/reasoning_benchmark/answer_extraction.py b/benchmark/reasoning_benchmark/answer_extraction.py new file mode 100644 index 000000000..45ce59afb --- /dev/null +++ b/benchmark/reasoning_benchmark/answer_extraction.py @@ -0,0 +1,269 @@ +# Adapted from https://github.com/deepseek-ai/DeepSeek-Math/blob/main/evaluation/data_processing/answer_extraction.py + +import re + +import regex + + +def _fix_fracs(string): + substrs = string.split("\\frac") + new_str = substrs[0] + if len(substrs) > 1: + substrs = substrs[1:] + for substr in substrs: + new_str += "\\frac" + if len(substr) > 0 and substr[0] == "{": + new_str += substr + else: + try: + assert len(substr) >= 2 + except: + return string + a = substr[0] + b = substr[1] + if b != "{": + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}{" + b + "}" + post_substr + else: + new_str += "{" + a + "}{" + b + "}" + else: + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}" + b + post_substr + else: + new_str += "{" + a + "}" + b + string = new_str + return string + + +def _fix_a_slash_b(string): + if len(string.split("/")) != 2: + return string + a = string.split("/")[0] + b = string.split("/")[1] + try: + if "sqrt" not in a: + a = int(a) + if "sqrt" not in b: + b = int(b) + assert string == "{}/{}".format(a, b) + new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" + return new_string + except: + return string + + +def _fix_sqrt(string): + _string = re.sub(r"\\sqrt(-?[0-9.a-zA-Z]+)", r"\\sqrt{\1}", string) + _string = re.sub(r"\\sqrt\s+(\w+)$", r"\\sqrt{\1}", _string) + return _string + + +def _fix_tan(string): + _string = re.sub(r"\\tan(-?[0-9.a-zA-Z]+)", r"\\tan{\1}", string) + _string = re.sub(r"\\tan\s+(\w+)$", r"\\tan{\1}", _string) + return _string + + +def strip_string(string): + string = str(string).strip() + # linebreaks + string = string.replace("\n", "") + + # right "." + string = string.rstrip(".") + + # remove inverse spaces + string = string.replace("\\!", "") + # string = string.replace("\\ ", "") + + # replace \\ with \ + # string = string.replace("\\\\", "\\") + # string = string.replace("\\\\", "\\") + + if string.startswith("\\text{") and string.endswith("}"): + string = string.split("{", 1)[1][:-1] + + # replace tfrac and dfrac with frac + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + string = string.replace("cfrac", "frac") + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + + # Remove unit: miles, dollars if after is not none + _string = re.sub(r"\\text{.*?}$", "", string).strip() + if _string != "" and _string != string: + # print("Warning: unit not removed: '{}' -> '{}'".format(string, _string)) + string = _string + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "").strip() + string = string.replace("^\\circ", "").strip() + + string = regex.sub(r"\{(c|m)?m\}(\^(2|3))?", "", string).strip() + string = regex.sub(r"p\.m\.$", "", string).strip() + string = regex.sub(r"(\d)\s*t$", r"\1", string).strip() + + # remove dollar signs + string = string.replace("\\$", "") + string = string.replace("$", "") + + # string = string.replace("\\text", "") + string = string.replace("x\\in", "") + + # remove percentage + string = string.replace("\\%", "%") + string = string.replace("\%", "%") + # string = string.replace("%", "") + + # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string + string = string.replace(" .", " 0.") + string = string.replace("{.", "{0.") + + # cdot + string = string.replace("\\cdot", "") + + # inf + string = string.replace("infinity", "\\infty") + if "\\infty" not in string: + string = string.replace("inf", "\\infty") + string = string.replace("+\\inity", "\\infty") + + # and + # string = string.replace("and", "") + string = string.replace("\\mathbf", "") + string = string.replace("\\mathrm", "") + + # use regex to remove \mbox{...} + string = re.sub(r"\\mbox{.*?}", "", string) + + # quote + string.replace("'", "") + string.replace('"', "") + + # i, j + if "j" in string and "i" not in string: + string = string.replace("j", "i") + + # replace a.000b where b is not number or b is end, with ab, use regex + string = re.sub(r"(\d+)\.0+([^\d])", r"\1\2", string) + string = re.sub(r"(\d+)\.0+$", r"\1", string) + + # if empty, return empty string + if len(string) == 0: + return string + if string[0] == ".": + string = "0" + string + + # to consider: get rid of e.g. "k = " or "q = " at beginning + # if len(string.split("=")) == 2: + # if len(string.split("=")[0]) <= 2: + # string = string.split("=")[1] + + string = _fix_sqrt(string) + string = _fix_tan(string) + string = string.replace(" ", "") + + # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} + string = _fix_fracs(string) + + # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y + string = _fix_a_slash_b(string) + + string = regex.sub(r"(\\|,|\.)+$", "", string) + + return string + + +def extract_boxed_answers(text): + answers = [] + for piece in text.split("boxed{")[1:]: + n = 0 + for i in range(len(piece)): + if piece[i] == "{": + n += 1 + elif piece[i] == "}": + n -= 1 + if n < 0: + if i + 1 < len(piece) and piece[i + 1] == "%": + answers.append(piece[: i + 1]) + else: + answers.append(piece[:i]) + break + return answers + + +def extract_program_output(pred_str): + """ + extract output between the last ```output\n...\n``` + """ + if "```output" not in pred_str: + return "" + if "```output" in pred_str: + pred_str = pred_str.split("```output")[-1] + if "```" in pred_str: + pred_str = pred_str.split("```")[0] + output = pred_str.strip() + return output + + +def extract_answer(pred_str, exhaust=False): + pred = [] + if "final answer is $" in pred_str and "$. I hope" in pred_str: + tmp = pred_str.split("final answer is $", 1)[1] + pred = [tmp.split("$. I hope", 1)[0].strip()] + elif "boxed" in pred_str: + pred = extract_boxed_answers(pred_str) + elif "he answer is" in pred_str: + pred = [pred_str.split("he answer is")[-1].strip()] + else: + program_output = extract_program_output(pred_str) + if program_output != "": + # fall back to program + pred.append(program_output) + else: # use the last number + pattern = "-?\d*\.?\d+" + ans = re.findall(pattern, pred_str.replace(",", "")) + if len(ans) >= 1: + ans = ans[-1] + else: + ans = "" + if ans: + pred.append(ans) + + # multiple line + _pred = [] + for ans in pred: + ans = ans.strip().split("\n")[0] + ans = ans.lstrip(":") + ans = ans.rstrip(".") + ans = ans.rstrip("/") + ans = strip_string(ans) + _pred.append(ans) + if exhaust: + return _pred + else: + return _pred[-1] if _pred else "" + + +def extract_math_answer(question, reasoning, task): + answer = [] + for ans in extract_answer(reasoning, exhaust=True): + if "separated by commas" in question and all(ch not in ans for ch in "()[]"): + answer.extend([a.strip() for a in ans.split(",")]) + elif regex.search(r"\\text\{\s*and\s*\}", ans): + answer.extend( + [ + a.strip() + for a in regex.sub(r"\\text\{\s*and\s*\}", "[SEP]", ans).split( + "[SEP]" + ) + ] + ) + else: + answer.append(ans.strip()) + return answer diff --git a/benchmark/reasoning_benchmark/bench_sglang.py b/benchmark/reasoning_benchmark/bench_sglang.py new file mode 100644 index 000000000..ccbff9d17 --- /dev/null +++ b/benchmark/reasoning_benchmark/bench_sglang.py @@ -0,0 +1,135 @@ +import argparse +import json +import time + +import answer_extraction +import eval_utils +import numpy as np +from datasets import load_dataset + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +@sgl.function +def reasoning_gen(s, question: str): + s += sgl.user( + question + + "\nPlease reason step by step, and put your final answer within \boxed{}." + ) + s += sgl.assistant( + sgl.gen( + "answer", + ) + ) + + +def convert_dataset(path: str, question_key: str, answer_key: str, num_tries: int): + raw_dataset = load_dataset(path) + questions = [] + answers = [] + for data in raw_dataset["train"]: + question = data[question_key] + answer = data[answer_key] + for _ in range(num_tries): + questions.append({"question": question}) + answers.append({"answer": answer}) + return questions, answers + + +def main(args): + # Select backend + sgl.set_default_backend(select_sglang_backend(args)) + + # Get dataset + questions, answers = convert_dataset( + args.data_path, args.question_key, args.answer_key, args.num_tries + ) + + # Run requests + tic = time.perf_counter() + states = reasoning_gen.run_batch( + questions, + num_threads=args.parallel, + progress_bar=True, + temperature=0.6, + max_new_tokens=32768, + top_p=0.95, + ) + latency = time.perf_counter() - tic + + # Extract results and record outcomes in a list. + outcomes = [] + for i, state in enumerate(states): + try: + pred_answer = answer_extraction.extract_math_answer( + questions[i]["question"], state["answer"], "limo" + ) + gt_answer = str(answers[i]["answer"]) + pred_answer = ( + pred_answer[-1] if isinstance(pred_answer, list) else pred_answer + ) + is_correct = 1 if eval_utils.math_equal(pred_answer, gt_answer) else 0 + except Exception as e: + print(f"Error extracting answer: {e}") + is_correct = 0 + + outcomes.append(is_correct) + + # Calculate overall accuracy using numpy + overall_accuracy = np.mean(outcomes) + print(f"Overall Accuracy: {overall_accuracy}") + + # Calculate mean standard error over questions if num_tries >= 2 + if args.num_tries > 1: + outcomes_np = np.array(outcomes).reshape(-1, args.num_tries) + # Using sample standard deviation with ddof=1 + std_per_question = np.std(outcomes_np, axis=1, ddof=1) + # Compute the standard error for each question: std / sqrt(num_tries) + se_per_question = std_per_question / np.sqrt(args.num_tries) + mean_se = se_per_question.mean() + print(f"Mean Standard Error of Accuracy across questions: {mean_se}") + else: + mean_se = None + print("Not enough samples per question to compute standard error.") + + # Calculate output throughput + num_output_tokens = sum( + s.get_meta_info("answer")["completion_tokens"] for s in states + ) + output_throughput = num_output_tokens / latency + print(f"Output throughput: {output_throughput} token/s") + + # Dump results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "limo", + "backend": args.backend, + "latency": round(latency, 3), + "overall_accuracy": round(overall_accuracy, 3), + "mean_se_accuracy": round(mean_se, 3) if mean_se is not None else None, + "num_requests": len(questions), + "other": { + "num_questions": len(questions), + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="GAIR/LIMO") + parser.add_argument("--question-key", type=str, default="question") + parser.add_argument("--answer-key", type=str, default="answer") + parser.add_argument("--num-tries", type=int, default=1) + add_common_sglang_args_and_parse(parser) + args = parser.parse_args() + main(args) diff --git a/benchmark/reasoning_benchmark/eval_utils.py b/benchmark/reasoning_benchmark/eval_utils.py new file mode 100644 index 000000000..ab736954f --- /dev/null +++ b/benchmark/reasoning_benchmark/eval_utils.py @@ -0,0 +1,206 @@ +# Adapted from https://github.com/deepseek-ai/DeepSeek-Math/blob/main/evaluation/eval/eval_utils.py + +from math import isclose + +import regex +from sympy import N, simplify +from sympy.parsing.latex import parse_latex +from sympy.parsing.sympy_parser import parse_expr + + +def parse_digits(num): + # format: 234.23 || 23% + num = regex.sub(",", "", str(num)) + try: + return float(num) + except: + if num.endswith("%"): + num = num[:-1] + if num.endswith("\\"): + num = num[:-1] + try: + return float(num) / 100 + except: + pass + return None + + +def is_digit(num): + # paired with parse_digits + return parse_digits(num) is not None + + +def symbolic_equal(a, b): + def _parse(s): + for f in [parse_latex, parse_expr]: + try: + return f(s) + except: + pass + return s + + a = _parse(a) + b = _parse(b) + + try: + if simplify(a - b) == 0: + return True + except: + pass + + try: + if isclose(N(a), N(b), abs_tol=1e-3): + return True + except: + pass + return False + + +def math_equal(prediction, reference, include_percentage=True, is_close=True): + """ + Exact match of math if and only if: + 1. numerical equal: both can convert to float and are equal + 2. symbolic equal: both can convert to sympy expression and are equal + """ + if str(prediction) == str(reference): + return True + + try: # 1. numerical equal + if is_digit(prediction) and is_digit(reference): + prediction = parse_digits(prediction) + reference = parse_digits(reference) + # number questions + if include_percentage: + gt_result = [reference / 100, reference, reference * 100] + else: + gt_result = [reference] + for item in gt_result: + try: + if is_close: + if isclose(item, prediction, abs_tol=1e-3): + return True + else: + if item == prediction: + return True + except Exception: + continue + return False + except: + pass + + if not prediction and prediction not in [0, False]: + return False + + # 2. symbolic equal + reference = str(reference).strip() + prediction = str(prediction).strip() + + if ( + regex.match(r"(\(|\[).+(\)|\])", prediction) is not None + and regex.match(r"(\(|\[).+(\)|\])", reference) is not None + ): + pred_parts = prediction[1:-1].split(",") + ref_parts = reference[1:-1].split(",") + if len(pred_parts) == len(ref_parts): + if all( + [ + math_equal( + pred_parts[i], ref_parts[i], include_percentage, is_close + ) + for i in range(len(pred_parts)) + ] + ): + return True + + # Add back matrix comparison + if ( + ( + prediction.startswith("\\begin{pmatrix}") + or prediction.startswith("\\begin{bmatrix}") + ) + and ( + prediction.endswith("\\end{pmatrix}") + or prediction.endswith("\\end{bmatrix}") + ) + and ( + reference.startswith("\\begin{pmatrix}") + or reference.startswith("\\begin{bmatrix}") + ) + and ( + reference.endswith("\\end{pmatrix}") or reference.endswith("\\end{bmatrix}") + ) + ): + pred_lines = [ + line.strip() + for line in prediction[ + len("\\begin{pmatrix}") : -len("\\end{pmatrix}") + ].split("\\\\") + if line.strip() + ] + ref_lines = [ + line.strip() + for line in reference[ + len("\\begin{pmatrix}") : -len("\\end{pmatrix}") + ].split("\\\\") + if line.strip() + ] + matched = True + if len(pred_lines) == len(ref_lines): + for pred_line, ref_line in zip(pred_lines, ref_lines): + pred_parts = pred_line.split("&") + ref_parts = ref_line.split("&") + if len(pred_parts) == len(ref_parts): + if not all( + [ + math_equal( + pred_parts[i], + ref_parts[i], + include_percentage, + is_close, + ) + for i in range(len(pred_parts)) + ] + ): + matched = False + break + else: + matched = False + if not matched: + break + else: + matched = False + if matched: + return True + + # Add back equation comparison + if prediction.count("=") == 1 and reference.count("=") == 1: + pred = prediction.split("=") + pred = f"{pred[0].strip()} - ({pred[1].strip()})" + ref = reference.split("=") + ref = f"{ref[0].strip()} - ({ref[1].strip()})" + if symbolic_equal(pred, ref) or symbolic_equal(f"-({pred})", ref): + return True + elif ( + prediction.count("=") == 1 + and len(prediction.split("=")[0].strip()) <= 2 + and "=" not in reference + ): + if math_equal( + prediction.split("=")[1], reference, include_percentage, is_close + ): + return True + elif ( + reference.count("=") == 1 + and len(reference.split("=")[0].strip()) <= 2 + and "=" not in prediction + ): + if math_equal( + prediction, reference.split("=")[1], include_percentage, is_close + ): + return True + + # symbolic equal with sympy + if symbolic_equal(prediction, reference): + return True + + return False diff --git a/benchmark/tip_suggestion/.gitignore b/benchmark/tip_suggestion/.gitignore new file mode 100644 index 000000000..3322de9b9 --- /dev/null +++ b/benchmark/tip_suggestion/.gitignore @@ -0,0 +1 @@ +!topic.jsonl diff --git a/benchmark/tip_suggestion/README.md b/benchmark/tip_suggestion/README.md new file mode 100644 index 000000000..97ebe1d82 --- /dev/null +++ b/benchmark/tip_suggestion/README.md @@ -0,0 +1,33 @@ +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 64 +python3 bench_sglang.py --num-questions 32 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --backend vllm --num-questions 64 +``` + + +### Benchmark guidance +``` +python3 bench_other.py --backend guidance --num-questions 32 --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --backend lmql --num-questions 32 --parallel 1 +``` diff --git a/benchmark/tip_suggestion/bench_other.py b/benchmark/tip_suggestion/bench_other.py new file mode 100644 index 000000000..2630081bd --- /dev/null +++ b/benchmark/tip_suggestion/bench_other.py @@ -0,0 +1,133 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +number = 5 + + +def expand_tip(topic, tip, generate): + s = ( + """Please expand a tip for a topic into a detailed paragraph. + +Topic: staying healthy +Tip: Regular Exercise +Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement. + +Topic: building a campfire +Tip: Choose the Right Location +Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches. + +Topic: writing a blog post +Tip: structure your content effectively +Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement. + +Topic: """ + + topic + + "\nTip: " + + tip + + "\nParagraph:" + ) + return generate(s, max_tokens=128, stop=["\n\n"]) + + +def suggest_tips(topic, generate): + s = "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n" + s += "USER: Give some tips for " + topic + ".\n" + s += ( + "ASSISTANT: Okay. Here are " + + str(number) + + " concise tips, each under 8 words:\n" + ) + + tips = [] + for i in range(1, 1 + number): + s += f"{i}." + tip = generate(s, max_tokens=24, stop=[".", "\n"]) + s += tip + ".\n" + tips.append(tip) + + paragraphs = [expand_tip(topic, tip, generate=generate) for tip in tips] + + for i in range(1, 1 + number): + s += f"Tip {i}:" + paragraphs[i - 1] + "\n" + return s + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + states = [None] * len(lines) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + tic = time.perf_counter() + if args.backend != "lmql": + + def get_one_answer(i): + states[i] = suggest_tips(lines[i]["topic"], call_generate) + + if args.parallel == 1: + for i in tqdm(range(len(lines))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(lines)))), + total=len(lines), + ) + ) + + else: + import asyncio + + from lmql_funcs import suggest_tips_async + + async def get_one_answer_async(i): + states[i] = await suggest_tips_async(lines[i]["topic"], call_generate) + + batches = [] + for i in range(0, len(lines), args.parallel): + batches.append(list(range(i, min(i + args.parallel, len(lines))))) + loop = asyncio.get_event_loop() + for batch in tqdm(batches): + loop.run_until_complete( + asyncio.gather(*[get_one_answer_async(i) for i in batch]) + ) + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "tip_suggestion", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="topic.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/tip_suggestion/bench_sglang.py b/benchmark/tip_suggestion/bench_sglang.py new file mode 100644 index 000000000..86c476f97 --- /dev/null +++ b/benchmark/tip_suggestion/bench_sglang.py @@ -0,0 +1,100 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +number = 5 + + +@sgl.function +def expand_tip(s, topic, tip): + s += ( + """Please expand a tip for a topic into a detailed paragraph. + +Topic: staying healthy +Tip: Regular Exercise +Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement. + +Topic: building a campfire +Tip: Choose the Right Location +Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches. + +Topic: writing a blog post +Tip: structure your content effectively +Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement. + +Topic: """ + + topic + + "\nTip: " + + tip + + "\nParagraph:" + ) + s += sgl.gen("paragraph", max_tokens=128, stop=["\n\n"], temperature=0) + + +@sgl.function +def suggest_tips(s, topic): + s += "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n" + s += "USER: Give some tips for " + topic + ".\n" + s += ( + "ASSISTANT: Okay. Here are " + + str(number) + + " concise tips, each under 8 words:\n" + ) + + paragraphs = [] + for i in range(1, 1 + number): + s += f"{i}." + sgl.gen(f"tip_{i}", max_tokens=24, stop=[".", "\n"]) + ".\n" + paragraphs.append(expand_tip(topic=topic, tip=s[f"tip_{i}"])) + + for i in range(1, 1 + number): + s += f"Tip {i}:" + paragraphs[i - 1]["paragraph"] + "\n" + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + arguments = [{"topic": l["topic"]} for l in lines] + + # Select backend + sgl.set_default_backend(select_sglang_backend(args)) + + # Run requests + tic = time.perf_counter() + states = suggest_tips.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.perf_counter() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "tip_suggestion", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="topic.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/tip_suggestion/lmql_funcs.py b/benchmark/tip_suggestion/lmql_funcs.py new file mode 100644 index 000000000..7790bbe95 --- /dev/null +++ b/benchmark/tip_suggestion/lmql_funcs.py @@ -0,0 +1,50 @@ +number = 5 + + +async def expand_tip_async(topic, tip, generate): + s = ( + """Please expand a tip for a topic into a detailed paragraph. + +Topic: staying healthy +Tip: Regular Exercise +Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement. + +Topic: building a campfire +Tip: Choose the Right Location +Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches. + +Topic: writing a blog post +Tip: structure your content effectively +Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement. + +Topic: """ + + topic + + "\nTip: " + + tip + + "\nParagraph:" + ) + return await generate(s, max_tokens=128, stop="\n\n") + + +async def suggest_tips_async(topic, generate): + s = "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n" + s += "USER: Give some tips for " + topic + ".\n" + s += ( + "ASSISTANT: Okay. Here are " + + str(number) + + " concise tips, each under 8 words:\n" + ) + + tips = [] + for i in range(1, 1 + number): + s += f"{i}." + # NOTE: stop is different due to lmql does not support a list of stop tokens + tip = await generate(s, max_tokens=24, stop=".\n") + s += tip + ".\n" + tips.append(tip) + + paragraphs = [await expand_tip_async(topic, tip, generate=generate) for tip in tips] + + for i in range(1, 1 + number): + s += f"Tip {i}:" + paragraphs[i - 1] + "\n" + return s diff --git a/benchmark/tip_suggestion/topic.jsonl b/benchmark/tip_suggestion/topic.jsonl new file mode 100644 index 000000000..c8ac6d12e --- /dev/null +++ b/benchmark/tip_suggestion/topic.jsonl @@ -0,0 +1,50 @@ +{"topic": "organizing a successful charity event", "number": 6} +{"topic": "improving personal credit scores", "number": 7} +{"topic": "staying motivated during job searches", "number": 5} +{"topic": "maintaining a work-life balance", "number": 9} +{"topic": "reducing carbon footprint at home", "number": 8} +{"topic": "starting a book club", "number": 5} +{"topic": "learning to play a musical instrument", "number": 7} +{"topic": "getting into freelance writing", "number": 6} +{"topic": "beginner yoga poses", "number": 8} +{"topic": "preparing for graduate school exams", "number": 5} +{"topic": "exploring minimalist living", "number": 9} +{"topic": "effective grocery shopping", "number": 7} +{"topic": "winter camping", "number": 5} +{"topic": "starting a podcast on a budget", "number": 8} +{"topic": "creating a capsule wardrobe", "number": 6} +{"topic": "improving your writing skills", "number": 7} +{"topic": "learning a new software quickly", "number": 9} +{"topic": "reducing anxiety before public speaking", "number": 5} +{"topic": "planning a solo travel adventure", "number": 8} +{"topic": "beginner skateboarders", "number": 6} +{"topic": "studying abroad", "number": 7} +{"topic": "planting a vegetable garden", "number": 5} +{"topic": "adopting a shelter pet", "number": 9} +{"topic": "learning to cook ethnic cuisines", "number": 8} +{"topic": "effective conflict resolution", "number": 5} +{"topic": "starting a vlog", "number": 7} +{"topic": "keeping a daily journal", "number": 6} +{"topic": "improving sleep hygiene", "number": 8} +{"topic": "beginner mountain climbers", "number": 5} +{"topic": "creating a mobile app", "number": 9} +{"topic": "maintaining a saltwater aquarium", "number": 7} +{"topic": "preparing for a baby's arrival", "number": 6} +{"topic": "writing a fantasy novel", "number": 5} +{"topic": "effective team leadership", "number": 8} +{"topic": "making a documentary film", "number": 9} +{"topic": "learning about historical events", "number": 7} +{"topic": "baking gluten-free treats", "number": 6} +{"topic": "improving mental arithmetic skills", "number": 5} +{"topic": "building a treehouse", "number": 8} +{"topic": "getting started with watercolor painting", "number": 9} +{"topic": "creating a YouTube tutorial series", "number": 7} +{"topic": "landscape photography", "number": 5} +{"topic": "navigating cultural differences", "number": 6} +{"topic": "preparing for a marathon", "number": 8} +{"topic": "building an online business", "number": 9} +{"topic": "learning to dance at home", "number": 5} +{"topic": "self-publishing a book", "number": 7} +{"topic": "starting an urban farm", "number": 6} +{"topic": "improving your memory", "number": 8} +{"topic": "creating a personal brand online", "number": 9} diff --git a/benchmark/tree_of_thought_deep/README.md b/benchmark/tree_of_thought_deep/README.md new file mode 100644 index 000000000..bf5ab1638 --- /dev/null +++ b/benchmark/tree_of_thought_deep/README.md @@ -0,0 +1,51 @@ +## Download data +``` +wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl +``` + +## Run benchmark + +NOTE: This is an implementation for throughput/latency benchmark purposes. The prompts are not tuned to achieve good accuracy on the GSM-8K tasks. + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 32 +python3 bench_sglang.py --num-questions 16 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 32 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 32 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 8 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --num-questions 8 --backend lmql --parallel 1 +``` diff --git a/benchmark/tree_of_thought_deep/bench_other.py b/benchmark/tree_of_thought_deep/bench_other.py new file mode 100644 index 000000000..0ef8c6360 --- /dev/null +++ b/benchmark/tree_of_thought_deep/bench_other.py @@ -0,0 +1,222 @@ +import argparse +import ast +import json +import re +import time +from collections import Counter +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def most_frequent_number(numbers): + if not numbers: + return None + + frequency = Counter(numbers) + most_frequent = max(frequency, key=frequency.get) + return most_frequent + + +USER_PREFIX = "[INST] " +USER_SUFFIX = " [/INST]" +ASSISTANT_PREFIX = "" +ASSISTANT_SUFFIX = " " + +# Use a low temp to make the results more deterministic and the comparison more fair. +temp = 0.001 + + +def propose_plan(s, question, num_branches, call_generate): + s += ( + USER_PREFIX + + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + + USER_SUFFIX + ) + + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def execute_plan(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def reflect_solution(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def get_final_answer(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def tree_search(question, num_branches, call_generate): + plan_forks = propose_plan("", question, num_branches, call_generate) + + sol_states = [] + for plan in plan_forks: + forks = execute_plan(plan, num_branches, call_generate) + sol_states.extend(forks) + + ref_states = [] + for sol in sol_states: + forks = reflect_solution(sol, num_branches, call_generate) + ref_states.extend(forks) + + solutions = [] + for sol in ref_states: + ans = get_final_answer(sol, num_branches, call_generate) + solutions.append(ans) + + return solutions + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + num_branches = 2 + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q, "num_branches": num_branches} for q in questions] + + # Select backend + call_generate = get_call_generate(args) + + # Run requests + states = [None] * len(questions) + + tic = time.perf_counter() + if args.backend != "lmql": + + def get_one_answer(i): + states[i] = tree_search(**arguments[i], call_generate=call_generate) + + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + else: + import asyncio + + from lmql_funcs import tree_search_async + + async def get_one_answer_async(i): + states[i] = await tree_search_async( + **arguments[i], call_generate=call_generate + ) + + batches = [ + [] for _ in range((len(questions) + args.parallel - 1) // args.parallel) + ] + for i in range(len(questions)): + batches[i // args.parallel].append(i) + + loop = asyncio.get_event_loop() + for bt in tqdm(batches): + tasks = [get_one_answer_async(k) for k in bt] + loop.run_until_complete(asyncio.gather(*tasks)) + + latency = time.perf_counter() - tic + + answers_text = [] + for s in states: + answers_text.append([x for xs in s for x in xs]) + + preds = [] + for i in range(len(states)): + answers = [get_answer_value(v) for v in answers_text[i]] + preds.append(most_frequent_number(answers)) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) + + with open(args.result_file, "a") as fout: + value = { + "task": "tree_of_thought_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/tree_of_thought_deep/bench_sglang.py b/benchmark/tree_of_thought_deep/bench_sglang.py new file mode 100644 index 000000000..bcdb6e54d --- /dev/null +++ b/benchmark/tree_of_thought_deep/bench_sglang.py @@ -0,0 +1,171 @@ +import argparse +import ast +import json +import re +import time +from collections import Counter + +import numpy as np + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def most_frequent_number(numbers): + if not numbers: + return None + + frequency = Counter(numbers) + most_frequent = max(frequency, key=frequency.get) + return most_frequent + + +# Use a low temp to make the results more deterministic and the comparison more fair. +temp = 0.001 + + +def propose_plan(s, question, num_branches): + s += sgl.user( + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("plan", max_tokens=256, temperature=temp)) + return forks + + +def execute_plan(s, num_branches): + s += sgl.user( + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("answer", max_tokens=256, temperature=temp)) + return forks + + +def reflect_solution(s, num_branches): + s += sgl.user( + """Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("score", max_tokens=256, temperature=temp)) + return forks + + +def get_final_answer(s, num_branches): + s += sgl.user( + """Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("final_answer", max_tokens=256, temperature=temp)) + return forks + + +@sgl.function +def tree_search(s, question, num_branches): + plan_forks = propose_plan(s, question, num_branches) + + sol_states = [] + for plan in plan_forks: + forks = execute_plan(plan, num_branches) + sol_states.extend(forks) + + ref_states = [] + for sol in sol_states: + forks = reflect_solution(sol, num_branches) + ref_states.extend(forks) + + solutions = [] + for sol in ref_states: + forks = get_final_answer(sol, num_branches) + solutions.append(forks) + solutions = [[s.text() for s in forks] for forks in solutions] + + return solutions + + +def main(args): + lines = read_jsonl(args.data_path) + lines = list(lines) + + # Construct prompts + num_branches = 2 + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q, "num_branches": num_branches} for q in questions] + + # Select backend + backend = select_sglang_backend(args) + + # Run requests + tic = time.perf_counter() + states = tree_search.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + answers_text = [] + for s in states: + answers_text.append([x for xs in s.ret_value for x in xs]) + + preds = [] + for i in range(len(states)): + answers = [get_answer_value(v) for v in answers_text[i]] + preds.append(most_frequent_number(answers)) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) + + with open(args.result_file, "a") as fout: + value = { + "task": "tree_of_thought_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/benchmark/tree_of_thought_deep/lmql_funcs.py b/benchmark/tree_of_thought_deep/lmql_funcs.py new file mode 100644 index 000000000..c783cdbe3 --- /dev/null +++ b/benchmark/tree_of_thought_deep/lmql_funcs.py @@ -0,0 +1,82 @@ +from bench_other import ( + ASSISTANT_PREFIX, + ASSISTANT_SUFFIX, + USER_PREFIX, + USER_SUFFIX, + temp, +) + + +async def propose_plan_async(s, question, num_branches, call_generate): + s += ( + USER_PREFIX + + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + + USER_SUFFIX + ) + + s += ASSISTANT_PREFIX + comps = await call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +async def execute_plan_async(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = await call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +async def reflect_solution_async(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = await call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +async def get_final_answer_async(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = await call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +async def tree_search_async(question, num_branches, call_generate): + plan_forks = await propose_plan_async("", question, num_branches, call_generate) + + sol_states = [] + for plan in plan_forks: + forks = await execute_plan_async(plan, num_branches, call_generate) + sol_states.extend(forks) + + ref_states = [] + for sol in sol_states: + forks = await reflect_solution_async(sol, num_branches, call_generate) + ref_states.extend(forks) + + solutions = [] + for sol in ref_states: + ans = await get_final_answer_async(sol, num_branches, call_generate) + solutions.append(ans) + + return solutions diff --git a/benchmark/tree_of_thought_v0/README.md b/benchmark/tree_of_thought_v0/README.md new file mode 100644 index 000000000..821bb20d1 --- /dev/null +++ b/benchmark/tree_of_thought_v0/README.md @@ -0,0 +1,43 @@ +## Download data +``` +wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl +``` + +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 32 --parallel 16 +python3 bench_sglang.py --num-questions 10 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 32 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 32 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 32 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` diff --git a/benchmark/tree_of_thought_v0/bench_other.py b/benchmark/tree_of_thought_v0/bench_other.py new file mode 100644 index 000000000..703ecd7f4 --- /dev/null +++ b/benchmark/tree_of_thought_v0/bench_other.py @@ -0,0 +1,179 @@ +import argparse +import ast +import json +import re +import time +from collections import Counter +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def most_frequent_number(numbers): + if not numbers: + return None + + frequency = Counter(numbers) + most_frequent = max(frequency, key=frequency.get) + return most_frequent + + +USER_PREFIX = "[INST] " +USER_SUFFIX = " [/INST]" +ASSISTANT_PREFIX = "" +ASSISTANT_SUFFIX = " " + +# Use a low temp to make the results more deterministic and the comparison more fair. +temp = 0.3 + + +def propose_plan(s, question, num_branches, call_generate): + s += ( + USER_PREFIX + + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + + USER_SUFFIX + ) + + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def execute_plan(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def reflect_solution(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Okay. Now you evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def tree_search(question, num_branches, call_generate): + s = "" + solutions = [] + + plan_forks = propose_plan(s, question, num_branches, call_generate) + for plan in plan_forks: + sol_forks = execute_plan(plan, num_branches, call_generate) + for sol in sol_forks: + score_forks = reflect_solution(sol, num_branches, call_generate) + solutions.append(sol_forks) + + return solutions + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + num_branches = 3 + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q, "num_branches": num_branches} for q in questions] + + # Select backend + call_generate = get_call_generate(args) + + # Run requests + states = [None] * len(questions) + + def get_one_answer(i): + states[i] = tree_search(**arguments[i], call_generate=call_generate) + + tic = time.perf_counter() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + latency = time.perf_counter() - tic + + answers_text = [] + for s in states: + answers_text.append([x for xs in s for x in xs]) + + preds = [] + for i in range(len(states)): + answers = [get_answer_value(v) for v in answers_text[i]] + preds.append(most_frequent_number(answers)) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) + + with open(args.result_file, "a") as fout: + value = { + "task": "tree_of_thought_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/benchmark/tree_of_thought_v0/bench_sglang.py b/benchmark/tree_of_thought_v0/bench_sglang.py new file mode 100644 index 000000000..6d7575f36 --- /dev/null +++ b/benchmark/tree_of_thought_v0/bench_sglang.py @@ -0,0 +1,159 @@ +import argparse +import ast +import json +import re +import time +from collections import Counter + +import numpy as np + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def most_frequent_number(numbers): + if not numbers: + return None + + frequency = Counter(numbers) + most_frequent = max(frequency, key=frequency.get) + return most_frequent + + +# Use a low temp to make the results more deterministic and the comparison more fair. +temp = 0.3 + + +def propose_plan(s, question, num_branches): + s += sgl.user( + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("plan", max_tokens=256, temperature=temp)) + return forks + + +def execute_plan(s, num_branches): + s += sgl.user( + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("answer", max_tokens=256, temperature=temp)) + return forks + + +def reflect_solution(s, num_branches): + s += sgl.user( + """Okay. Now you evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("score", max_tokens=256, temperature=temp)) + return forks + + +@sgl.function +def tree_search(s, question, num_branches): + forks_to_join = [] + + plan_forks = propose_plan(s, question, num_branches) + forks_to_join.append(plan_forks) + + sol_states = [] + for plan in plan_forks: + forks = execute_plan(plan, num_branches) + forks_to_join.append(forks) + sol_states.extend(forks) + + for sol in sol_states: + forks = reflect_solution(sol, num_branches) + forks_to_join.append(forks) + + for f in reversed(forks_to_join): + f.join() + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + num_branches = 3 + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q, "num_branches": num_branches} for q in questions] + + # Select backend + backend = select_sglang_backend(args) + + # Run requests + tic = time.perf_counter() + states = tree_search.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.perf_counter() - tic + answers_text = [] + for s in states: + answers_text.append([x for xs in s["answer"] for x in xs]) + + preds = [] + for i in range(len(states)): + answers = [get_answer_value(v) for v in answers_text[i]] + preds.append(most_frequent_number(answers)) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) + + with open(args.result_file, "a") as fout: + value = { + "task": "tree_of_thought_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 000000000..3b9a420b3 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,388 @@ +ARG CUDA_VERSION=12.9.1 +FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu22.04 as base + +ARG BUILD_TYPE=all +ARG BRANCH_TYPE=remote +ARG DEEPEP_COMMIT=b92d0d4860ce6866cd6d31bfbae937f9a7a3772b +ARG CMAKE_BUILD_PARALLEL_LEVEL=2 +ENV DEBIAN_FRONTEND=noninteractive \ + CUDA_HOME=/usr/local/cuda \ + GDRCOPY_HOME=/usr/src/gdrdrv-2.4.4/ \ + NVSHMEM_DIR=/sgl-workspace/nvshmem/install +# Add GKE default lib and bin locations. +ENV PATH="${PATH}:/usr/local/nvidia/bin" \ + LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/nvidia/lib:/usr/local/nvidia/lib64" + +RUN apt update && apt install wget -y && apt install software-properties-common -y \ + && add-apt-repository ppa:deadsnakes/ppa -y \ + && apt install python3.12-full python3.12-dev python3.10-venv -y \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 2 \ + && update-alternatives --set python3 /usr/bin/python3.12 \ + && wget https://bootstrap.pypa.io/get-pip.py \ + && python3 get-pip.py + +# Set timezone and install all packages +RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ + && apt-get update && apt-get install -y --no-install-recommends \ + tzdata \ + software-properties-common netcat-openbsd kmod unzip openssh-server \ + curl wget lsof zsh ccache tmux htop git-lfs tree \ + build-essential cmake \ + libopenmpi-dev libnuma1 libnuma-dev \ + libibverbs-dev libibverbs1 libibumad3 \ + librdmacm1 libnl-3-200 libnl-route-3-200 libnl-route-3-dev libnl-3-dev \ + ibverbs-providers infiniband-diags perftest \ + libgoogle-glog-dev libgtest-dev libjsoncpp-dev libunwind-dev \ + libboost-all-dev libssl-dev \ + libgrpc-dev libgrpc++-dev libprotobuf-dev protobuf-compiler protobuf-compiler-grpc \ + pybind11-dev \ + libhiredis-dev libcurl4-openssl-dev \ + libczmq4 libczmq-dev \ + libfabric-dev \ + patchelf \ + nvidia-dkms-550 \ + devscripts debhelper fakeroot dkms check libsubunit0 libsubunit-dev \ + && ln -sf /usr/bin/python3.12 /usr/bin/python \ + && rm -rf /var/lib/apt/lists/* \ + && apt-get clean + +# GDRCopy installation +RUN mkdir -p /tmp/gdrcopy && cd /tmp \ + && git clone https://github.com/NVIDIA/gdrcopy.git -b v2.4.4 \ + && cd gdrcopy/packages \ + && CUDA=/usr/local/cuda ./build-deb-packages.sh \ + && dpkg -i gdrdrv-dkms_*.deb libgdrapi_*.deb gdrcopy-tests_*.deb gdrcopy_*.deb \ + && cd / && rm -rf /tmp/gdrcopy + +# Fix DeepEP IBGDA symlink +RUN ln -sf /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so + +FROM scratch AS local_src +COPY . /src + +FROM base AS build-image +# Install SGLang +WORKDIR /sgl-workspace +ARG BRANCH_TYPE +COPY --from=local_src /src /tmp/local_src +RUN if [ "$BRANCH_TYPE" = "local" ]; then \ + cp -r /tmp/local_src /sgl-workspace/sglang; \ + else \ + git clone --depth=1 https://github.com/sgl-project/sglang.git /sgl-workspace/sglang; \ + fi \ + && rm -rf /tmp/local_src +RUN python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel html5lib six \ + && cd sglang \ + && case "$CUDA_VERSION" in \ + 12.6.1) CUINDEX=126 ;; \ + 12.8.1) CUINDEX=128 ;; \ + 12.9.1) CUINDEX=129 ;; \ + *) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \ + esac \ + && python3 -m pip install --no-cache-dir -e "python[${BUILD_TYPE}]" --extra-index-url https://download.pytorch.org/whl/cu${CUINDEX} \ + && python3 -m pip install --no-cache-dir nvidia-nccl-cu12==2.27.6 --force-reinstall --no-deps \ + && python3 -m flashinfer --download-cubin \ + && if [ "$CUDA_VERSION" = "12.6.1" ]; then \ + python3 -m pip install --no-cache-dir https://github.com/sgl-project/whl/releases/download/v0.3.9.post2/sgl_kernel-0.3.9.post2+cu124-cp310-abi3-manylinux2014_x86_64.whl --force-reinstall --no-deps ; \ + fi + +# Download source files +RUN wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.3.9/source/nvshmem_src_cuda12-all-all-3.3.9.tar.gz && \ + git clone https://github.com/deepseek-ai/DeepEP.git && \ + cd DeepEP && git checkout ${DEEPEP_COMMIT} && sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \ + cd .. && \ + tar -xf nvshmem_src_cuda12-all-all-3.3.9.tar.gz && \ + mv nvshmem_src nvshmem && \ + rm -f /sgl-workspace/nvshmem_src_cuda12-all-all-3.3.9.tar.gz + +# Build and install NVSHMEM +RUN cd /sgl-workspace/nvshmem && \ + NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -S . -B build/ -DCMAKE_INSTALL_PREFIX=${NVSHMEM_DIR} -DCMAKE_CUDA_ARCHITECTURES="90" && \ + cmake --build build --target install -j${CMAKE_BUILD_PARALLEL_LEVEL} + +# Install DeepEP +RUN cd /sgl-workspace/DeepEP && \ + case "$CUDA_VERSION" in \ + 12.6.1) \ + CHOSEN_TORCH_CUDA_ARCH_LIST='9.0' \ + ;; \ + 12.8.1|12.9.1) \ + CHOSEN_TORCH_CUDA_ARCH_LIST='9.0;10.0' \ + ;; \ + *) \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 \ + ;; \ + esac && \ + NVSHMEM_DIR=${NVSHMEM_DIR} TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" pip install . + +# Python tools +RUN python3 -m pip install --no-cache-dir \ + datamodel_code_generator \ + mooncake-transfer-engine==0.3.5 \ + pre-commit \ + pytest \ + black \ + isort \ + icdiff \ + uv \ + wheel \ + scikit-build-core \ + nixl \ + py-spy + +# Install development tools and utilities +RUN apt-get update && apt-get install -y \ + gdb \ + ninja-build \ + vim \ + tmux \ + htop \ + wget \ + curl \ + locales \ + lsof \ + git \ + git-lfs \ + zsh \ + tree \ + silversearcher-ag \ + cloc \ + unzip \ + pkg-config \ + libssl-dev \ + bear \ + ccache \ + less \ + && apt install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && rm -rf /var/lib/apt/lists/* \ + && apt-get clean + +RUN apt update -y \ + && apt install -y --no-install-recommends gnupg \ + && echo "deb http://developer.download.nvidia.com/devtools/repos/ubuntu2004/amd64 /" | tee /etc/apt/sources.list.d/nvidia-devtools.list \ + && apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub \ + && apt update -y \ + && apt install nsight-systems-cli -y + +# Set up locale +RUN locale-gen en_US.UTF-8 +ENV LANG en_US.UTF-8 +ENV LANGUAGE en_US:en +ENV LC_ALL en_US.UTF-8 + +# Install minimal Python packages +RUN python3 -m pip install --no-cache-dir --break-system-packages \ + pytest \ + black \ + isort \ + icdiff \ + scikit_build_core \ + uv \ + pre-commit \ + pandas \ + matplotlib \ + tabulate + +# Install diff-so-fancy +RUN curl -LSso /usr/local/bin/diff-so-fancy https://github.com/so-fancy/diff-so-fancy/releases/download/v1.4.4/diff-so-fancy \ + && chmod +x /usr/local/bin/diff-so-fancy + +# Install clang-format +RUN curl -LSso /usr/local/bin/clang-format https://github.com/muttleyxd/clang-tools-static-binaries/releases/download/master-32d3ac78/clang-format-16_linux-amd64 \ + && chmod +x /usr/local/bin/clang-format + +# Install clangd +RUN curl -L https://github.com/clangd/clangd/releases/download/18.1.3/clangd-linux-18.1.3.zip -o clangd.zip \ + && unzip clangd.zip \ + && cp -r clangd_18.1.3/bin/* /usr/local/bin/ \ + && cp -r clangd_18.1.3/lib/* /usr/local/lib/ \ + && rm -rf clangd_18.1.3 clangd.zip + +# Install CMake +RUN wget https://github.com/Kitware/CMake/releases/download/v3.31.1/cmake-3.31.1-linux-x86_64.tar.gz \ + && tar -xzf cmake-3.31.1-linux-x86_64.tar.gz \ + && cp -r cmake-3.31.1-linux-x86_64/bin/* /usr/local/bin/ \ + && cp -r cmake-3.31.1-linux-x86_64/share/* /usr/local/share/ \ + && rm -rf cmake-3.31.1-linux-x86_64 cmake-3.31.1-linux-x86_64.tar.gz + +# Install Rust toolchain for sgl-router +ENV PATH="/root/.cargo/bin:${PATH}" +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ + && rustc --version && cargo --version + +# Build and install sgl-router +RUN python3 -m pip install --no-cache-dir setuptools-rust \ + && cd /sgl-workspace/sglang/sgl-router \ + && cargo build --release \ + && python3 -m pip install --no-cache-dir . \ + && rm -rf /root/.cache + + +# Add yank script +COPY --chown=root:root <<-"EOF" /usr/local/bin/yank +#!/bin/bash +put() { + esc=$1 + test -n "$TMUX" -o -z "${TERM##screen*}" && esc="\033Ptmux;\033$esc\033\\" + printf "$esc" +} +put "\033]52;c;!\a" +buf=$( cat "$@" ) +len=$( printf %s "$buf" | wc -c ) max=74994 +test $len -gt $max && echo "$0: input is $(( len - max )) bytes too long" >&2 +put "\033]52;c;$( printf %s "$buf" | head -c $max | base64 | tr -d '\r\n' )\a" +test -n "$TMUX" && tmux set-buffer "$buf" ||: +EOF + +RUN chmod +x /usr/local/bin/yank + +# Install oh-my-zsh and plugins +RUN sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended \ + && git clone https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions \ + && git clone https://github.com/zsh-users/zsh-syntax-highlighting.git ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-syntax-highlighting + +# Configure Vim +COPY --chown=root:root <<-"EOF" /root/.vimrc +function! Yank(text) abort + let escape = system('yank', a:text) + if v:shell_error + echoerr escape + else + call writefile([escape], '/dev/tty', 'b') + endif +endfunction + +noremap y y:call Yank(@0) + +" automatically run yank(1) whenever yanking in Vim +function! CopyYank() abort + call Yank(join(v:event.regcontents, "\n")) +endfunction + +autocmd TextYankPost * call CopyYank() + +" Basic settings +set number +syntax on +set mouse=a +filetype indent on + +" Indentation +set autoindent nosmartindent +set smarttab +set expandtab +set shiftwidth=4 +set softtabstop=4 + +" Visual guides +set colorcolumn=120 +highlight ColorColumn ctermbg=5 + +" Status line +set laststatus=2 +set statusline=%<%f\ %h%m%r%=%{\"[\".(&fenc==\"\"?&enc:&fenc).((exists(\"+bomb\")\ &&\ &bomb)?\",B\":\"\").\"]\ \"}%k\ %-14.(%l,%c%V%)\ %P + +" Backspace behavior +set backspace=2 + +" Encoding +set encoding=utf-8 +set fileencoding=utf-8 +EOF + +# Configure tmux +COPY --chown=root:root <<-"EOF" /root/.tmux.conf +# Pane border styling +set -g pane-border-style fg='#742727',bg=black +set -g pane-active-border-style fg=red,bg=black + +# Status bar styling +set -g status-style bg='#0C8A92',fg=black + +# Change prefix key to backtick +set-option -g prefix ` +unbind C-b +bind-key ` send-prefix + +# Split panes using - and = with current path +unbind '"' +bind - splitw -v -c '#{pane_current_path}' +unbind '%' +bind = splitw -h -c '#{pane_current_path}' + +# Vi mode settings +bind-key -T copy-mode-vi Y send-keys -X copy-pipe 'yank > #{pane_tty}' +set-window-option -g mode-keys vi + +# Other settings +set-option -g escape-time 0 +set-option -g base-index 1 +set-window-option -g mouse on +set -g history-limit 100000 +EOF + +# Configure Git +RUN git config --global core.editor "vim" \ + && git config --global core.whitespace "fix,-indent-with-non-tab,trailing-space,cr-at-eol" \ + && git config --global core.pager "diff-so-fancy | less --tabs=4 -RFX" \ + && git config --global color.ui true \ + && git config --global color."diff-highlight".oldNormal "red bold" \ + && git config --global color."diff-highlight".oldHighlight "red bold 52" \ + && git config --global color."diff-highlight".newNormal "green bold" \ + && git config --global color."diff-highlight".newHighlight "green bold 22" \ + && git config --global color.diff.meta "11" \ + && git config --global color.diff.frag "magenta bold" \ + && git config --global color.diff.commit "yellow bold" \ + && git config --global color.diff.old "red bold" \ + && git config --global color.diff.new "green bold" \ + && git config --global color.diff.whitespace "red reverse" \ + && git config --global alias.lg "log --color --graph --pretty=format:'%Cred%h%Creset - %s %Cgreen(%cr) %C(bold blue)<%an>%Creset%C(auto)%d%Creset' --abbrev-commit --" \ + && git config --global http.sslVerify false \ + && git config --global pull.rebase true + +# Configure zsh +COPY --chown=root:root <<-"EOF" /root/.zshrc +export ZSH="/root/.oh-my-zsh" + +# Theme +ZSH_THEME="robbyrussell" + +# Plugins +plugins=( + git + z + zsh-autosuggestions + zsh-syntax-highlighting +) + +source $ZSH/oh-my-zsh.sh + +# Aliases +alias ll='ls -alF' +alias la='ls -A' +alias l='ls -CF' +alias vi='vim' + +# Enhanced history +HISTSIZE=10000 +SAVEHIST=10000 +setopt HIST_IGNORE_ALL_DUPS +setopt HIST_FIND_NO_DUPS +setopt INC_APPEND_HISTORY +EOF + +RUN set -euxo ; \ + curl --proto '=https' --tlsv1.2 -sSf https://just.systems/install.sh | bash -s -- --to /usr/local/bin + +# Set workspace directory +WORKDIR /sgl-workspace/sglang diff --git a/docker/Dockerfile.gb200 b/docker/Dockerfile.gb200 new file mode 100644 index 000000000..164326e23 --- /dev/null +++ b/docker/Dockerfile.gb200 @@ -0,0 +1,348 @@ +ARG CUDA_VERSION=12.9.1 +FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu22.04 + +ARG BUILD_TYPE=blackwell +ARG DEEPEP_COMMIT=1b14ad661c7640137fcfe93cccb2694ede1220b0 +ARG CMAKE_BUILD_PARALLEL_LEVEL=2 +ARG SGL_KERNEL_VERSION=0.3.9.post2 +ENV DEBIAN_FRONTEND=noninteractive \ + CUDA_HOME=/usr/local/cuda \ + GDRCOPY_HOME=/usr/src/gdrdrv-2.4.4/ \ + NVSHMEM_DIR=/sgl-workspace/nvshmem/install \ + BUILD_TYPE=${BUILD_TYPE} \ + TORCH_CUDA_ARCH_LIST="10.0 12.0" + +# Set timezone and install all packages +RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ + && apt-get update && apt-get install -y --no-install-recommends \ + tzdata \ + software-properties-common netcat-openbsd kmod unzip openssh-server \ + curl wget lsof zsh ccache tmux htop git-lfs tree \ + python3 python3-pip python3-dev libpython3-dev python3-venv \ + build-essential cmake \ + libopenmpi-dev libnuma1 libnuma-dev \ + libibverbs-dev libibverbs1 libibumad3 \ + librdmacm1 libnl-3-200 libnl-route-3-200 libnl-route-3-dev libnl-3-dev \ + ibverbs-providers infiniband-diags perftest \ + libgoogle-glog-dev libgtest-dev libjsoncpp-dev libunwind-dev \ + libboost-all-dev libssl-dev \ + libgrpc-dev libgrpc++-dev libprotobuf-dev protobuf-compiler-grpc \ + pybind11-dev \ + libhiredis-dev libcurl4-openssl-dev \ + libczmq4 libczmq-dev \ + libfabric-dev \ + patchelf \ + nvidia-dkms-550 \ + devscripts debhelper fakeroot dkms check libsubunit0 libsubunit-dev \ + && ln -sf /usr/bin/python3 /usr/bin/python \ + && rm -rf /var/lib/apt/lists/* \ + && apt-get clean + +# Install SGLang missing package for blackwell build type +RUN python3 -m pip install openai httpx + +# GDRCopy installation +RUN mkdir -p /tmp/gdrcopy && cd /tmp \ + && git clone https://github.com/NVIDIA/gdrcopy.git -b v2.4.4 \ + && cd gdrcopy/packages \ + && CUDA=/usr/local/cuda ./build-deb-packages.sh \ + && dpkg -i gdrdrv-dkms_*.deb libgdrapi_*.deb gdrcopy-tests_*.deb gdrcopy_*.deb \ + && cd / && rm -rf /tmp/gdrcopy + +# Fix DeepEP IBGDA symlink +RUN ln -sf /usr/lib/$(uname -m)-linux-gnu/libmlx5.so.1 /usr/lib/$(uname -m)-linux-gnu/libmlx5.so + +# Clone and install SGLang +WORKDIR /sgl-workspace +RUN python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel html5lib six \ + && git clone --depth 1 https://github.com/sgl-project/sglang.git \ + && cd sglang \ + && case "$CUDA_VERSION" in \ + 12.9.1) CUINDEX=129 ;; \ + *) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \ + esac \ + && if [ "$CUDA_VERSION" = "12.9.1" ]; then \ + python3 -m pip install --no-cache-dir nvidia-nccl-cu12==2.27.6 --force-reinstall --no-deps ; \ + python3 -m pip install --no-cache-dir https://github.com/sgl-project/whl/releases/download/v${SGL_KERNEL_VERSION}/sgl_kernel-${SGL_KERNEL_VERSION}+cu129-cp310-abi3-manylinux2014_$(uname -m).whl --force-reinstall --no-deps ; \ + fi \ + && python3 -m pip install --no-cache-dir -e "python[${BUILD_TYPE}]" --extra-index-url https://download.pytorch.org/whl/cu${CUINDEX} \ + && python3 -m flashinfer --download-cubin + +# Download source files +RUN wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.3.9/source/nvshmem_src_cuda12-all-all-3.3.9.tar.gz && \ + git clone https://github.com/fzyzcjy/DeepEP.git && \ + cd DeepEP && git checkout ${DEEPEP_COMMIT} && cd .. && \ + tar -xf nvshmem_src_cuda12-all-all-3.3.9.tar.gz && \ + mv nvshmem_src nvshmem && \ + rm -f /sgl-workspace/nvshmem_src_cuda12-all-all-3.3.9.tar.gz + +# Build and install NVSHMEM +RUN cd /sgl-workspace/nvshmem && \ + NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -S . -B build/ -DCMAKE_INSTALL_PREFIX=${NVSHMEM_DIR} -DCMAKE_CUDA_ARCHITECTURES="90;100;120" && \ + cmake --build build --target install -j${CMAKE_BUILD_PARALLEL_LEVEL} + +# Install DeepEP +RUN cd /sgl-workspace/DeepEP && \ + NVSHMEM_DIR=${NVSHMEM_DIR} pip install . + +# Python tools +RUN python3 -m pip install --no-cache-dir \ + datamodel_code_generator \ + mooncake-transfer-engine==0.3.5 \ + pre-commit \ + pytest \ + black \ + isort \ + icdiff \ + uv \ + wheel \ + scikit-build-core + +# Install nixl kv transfer backend +RUN python3 -m pip install --no-cache-dir \ + nixl + +# Install development tools and utilities +RUN apt-get update && apt-get install -y \ + gdb \ + ninja-build \ + vim \ + tmux \ + htop \ + wget \ + curl \ + locales \ + lsof \ + git \ + git-lfs \ + zsh \ + tree \ + silversearcher-ag \ + cloc \ + unzip \ + pkg-config \ + libssl-dev \ + bear \ + ccache \ + less \ + && apt install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && rm -rf /var/lib/apt/lists/* \ + && apt-get clean + +RUN apt update -y \ + && apt install -y --no-install-recommends gnupg \ + && echo "deb http://developer.download.nvidia.com/devtools/repos/ubuntu2004/$(if [ "$(uname -m)" = "aarch64" ]; then echo "arm64"; else echo "amd64"; fi) /" | tee /etc/apt/sources.list.d/nvidia-devtools.list \ + && apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/$(if [ "$(uname -m)" = "aarch64" ]; then echo "arm64"; else echo "x86_64"; fi)/7fa2af80.pub \ + && apt update -y \ + && apt install nsight-systems-cli -y + +# Set up locale +RUN locale-gen en_US.UTF-8 +ENV LANG=en_US.UTF-8 +ENV LANGUAGE=en_US:en +ENV LC_ALL=en_US.UTF-8 + +# Install minimal Python packages +RUN python3 -m pip install --no-cache-dir --break-system-packages \ + pytest \ + black \ + isort \ + icdiff \ + scikit_build_core \ + uv \ + pre-commit \ + pandas \ + matplotlib \ + tabulate + +# Install diff-so-fancy +RUN curl -LSso /usr/local/bin/diff-so-fancy https://github.com/so-fancy/diff-so-fancy/releases/download/v1.4.4/diff-so-fancy \ + && chmod +x /usr/local/bin/diff-so-fancy + +# Install clang-format +RUN curl -LSso /usr/local/bin/clang-format https://github.com/muttleyxd/clang-tools-static-binaries/releases/download/master-32d3ac78/clang-format-16_linux-amd64 \ + && chmod +x /usr/local/bin/clang-format + +# Install clangd +RUN curl -L https://github.com/clangd/clangd/releases/download/18.1.3/clangd-linux-18.1.3.zip -o clangd.zip \ + && unzip clangd.zip \ + && cp -r clangd_18.1.3/bin/* /usr/local/bin/ \ + && cp -r clangd_18.1.3/lib/* /usr/local/lib/ \ + && rm -rf clangd_18.1.3 clangd.zip + +# Install CMake +RUN CMAKE_VERSION=3.31.1 \ + && ARCH=$(uname -m) \ + && CMAKE_INSTALLER="cmake-${CMAKE_VERSION}-linux-${ARCH}" \ + && wget "https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/${CMAKE_INSTALLER}.tar.gz" \ + && tar -xzf "${CMAKE_INSTALLER}.tar.gz" \ + && cp -r "${CMAKE_INSTALLER}/bin/"* /usr/local/bin/ \ + && cp -r "${CMAKE_INSTALLER}/share/"* /usr/local/share/ \ + && rm -rf "${CMAKE_INSTALLER}" "${CMAKE_INSTALLER}.tar.gz" + +# Add yank script +COPY --chown=root:root <<-"EOF" /usr/local/bin/yank +#!/bin/bash +put() { + esc=$1 + test -n "$TMUX" -o -z "${TERM##screen*}" && esc="\033Ptmux;\033$esc\033\\" + printf "$esc" +} +put "\033]52;c;!\a" +buf=$( cat "$@" ) +len=$( printf %s "$buf" | wc -c ) max=74994 +test $len -gt $max && echo "$0: input is $(( len - max )) bytes too long" >&2 +put "\033]52;c;$( printf %s "$buf" | head -c $max | base64 | tr -d '\r\n' )\a" +test -n "$TMUX" && tmux set-buffer "$buf" ||: +EOF + +RUN chmod +x /usr/local/bin/yank + +# Install oh-my-zsh and plugins +RUN sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended \ + && git clone https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions \ + && git clone https://github.com/zsh-users/zsh-syntax-highlighting.git ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-syntax-highlighting + +# Configure Vim +COPY --chown=root:root <<-"EOF" /root/.vimrc +function! Yank(text) abort + let escape = system('yank', a:text) + if v:shell_error + echoerr escape + else + call writefile([escape], '/dev/tty', 'b') + endif +endfunction + +noremap y y:call Yank(@0) + +" automatically run yank(1) whenever yanking in Vim +function! CopyYank() abort + call Yank(join(v:event.regcontents, "\n")) +endfunction + +autocmd TextYankPost * call CopyYank() + +" Basic settings +set number +syntax on +set mouse=a +filetype indent on + +" Indentation +set autoindent nosmartindent +set smarttab +set expandtab +set shiftwidth=4 +set softtabstop=4 + +" Visual guides +set colorcolumn=120 +highlight ColorColumn ctermbg=5 + +" Status line +set laststatus=2 +set statusline=%<%f\ %h%m%r%=%{\"[\".(&fenc==\"\"?&enc:&fenc).((exists(\"+bomb\")\ &&\ &bomb)?\",B\":\"\").\"]\ \"}%k\ %-14.(%l,%c%V%)\ %P + +" Backspace behavior +set backspace=2 + +" Encoding +set encoding=utf-8 +set fileencoding=utf-8 +EOF + +# Configure tmux +COPY --chown=root:root <<-"EOF" /root/.tmux.conf +# Pane border styling +set -g pane-border-style fg='#742727',bg=black +set -g pane-active-border-style fg=red,bg=black + +# Status bar styling +set -g status-style bg='#0C8A92',fg=black + +# Change prefix key to backtick +set-option -g prefix ` +unbind C-b +bind-key ` send-prefix + +# Split panes using - and = with current path +unbind '"' +bind - splitw -v -c '#{pane_current_path}' +unbind '%' +bind = splitw -h -c '#{pane_current_path}' + +# Vi mode settings +bind-key -T copy-mode-vi Y send-keys -X copy-pipe 'yank > #{pane_tty}' +set-window-option -g mode-keys vi + +# Other settings +set-option -g escape-time 0 +set-option -g base-index 1 +set-window-option -g mouse on +EOF + +# Configure Git +RUN git config --global core.editor "vim" \ + && git config --global core.whitespace "fix,-indent-with-non-tab,trailing-space,cr-at-eol" \ + && git config --global core.pager "diff-so-fancy | less --tabs=4 -RFX" \ + && git config --global color.ui true \ + && git config --global color."diff-highlight".oldNormal "red bold" \ + && git config --global color."diff-highlight".oldHighlight "red bold 52" \ + && git config --global color."diff-highlight".newNormal "green bold" \ + && git config --global color."diff-highlight".newHighlight "green bold 22" \ + && git config --global color.diff.meta "11" \ + && git config --global color.diff.frag "magenta bold" \ + && git config --global color.diff.commit "yellow bold" \ + && git config --global color.diff.old "red bold" \ + && git config --global color.diff.new "green bold" \ + && git config --global color.diff.whitespace "red reverse" \ + && git config --global alias.lg "log --color --graph --pretty=format:'%Cred%h%Creset - %s %Cgreen(%cr) %C(bold blue)<%an>%Creset%C(auto)%d%Creset' --abbrev-commit --" \ + && git config --global http.sslVerify false \ + && git config --global pull.rebase true + +# Configure zsh +COPY --chown=root:root <<-"EOF" /root/.zshrc +export ZSH="/root/.oh-my-zsh" + +# Theme +ZSH_THEME="robbyrussell" + +# Plugins +plugins=( + git + z + zsh-autosuggestions + zsh-syntax-highlighting +) + +source $ZSH/oh-my-zsh.sh + +# Aliases +alias ll='ls -alF' +alias la='ls -A' +alias l='ls -CF' +alias vi='vim' + +# Enhanced history +HISTSIZE=10000 +SAVEHIST=10000 +setopt HIST_IGNORE_ALL_DUPS +setopt HIST_FIND_NO_DUPS +setopt INC_APPEND_HISTORY +EOF + +RUN set -euxo ; \ + curl --proto '=https' --tlsv1.2 -sSf https://just.systems/install.sh | bash -s -- --to /usr/local/bin + +# Set workspace directory +WORKDIR /sgl-workspace/sglang diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu new file mode 100644 index 000000000..3f9b0ae42 --- /dev/null +++ b/docker/Dockerfile.npu @@ -0,0 +1,93 @@ +ARG CANN_VERSION=8.2.rc1 +ARG DEVICE_TYPE=a3 +ARG OS=ubuntu22.04 +ARG PYTHON_VERSION=py3.11 + +FROM quay.io/ascend/cann:$CANN_VERSION-$DEVICE_TYPE-$OS-$PYTHON_VERSION + +# Update pip & apt sources +ARG PIP_INDEX_URL="https://pypi.org/simple/" +ARG APTMIRROR="" +ARG MEMFABRIC_URL=https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/mf_adapter-1.0.0-cp311-cp311-linux_aarch64.whl +ARG PYTORCH_VERSION=2.6.0 +ARG TORCHVISION_VERSION=0.21.0 +ARG PTA_URL="https://gitee.com/ascend/pytorch/releases/download/v7.1.0.1-pytorch2.6.0/torch_npu-2.6.0.post1-cp311-cp311-manylinux_2_28_aarch64.whl" +ARG VLLM_TAG=v0.8.5 +ARG TRITON_ASCEND_URL=https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/triton_ascend-3.2.0.dev20250729-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl +ARG SGLANG_TAG=main +ARG ASCEND_CANN_PATH=/usr/local/Ascend/ascend-toolkit +ARG SGLANG_KERNEL_NPU_TAG=main + +WORKDIR /workspace + +# Define environments +ENV DEBIAN_FRONTEND=noninteractive + +RUN pip config set global.index-url $PIP_INDEX_URL +RUN if [ -n "$APTMIRROR" ];then sed -i "s|.*.ubuntu.com|$APTMIRROR|g" /etc/apt/sources.list ;fi + +# Install development tools and utilities +RUN apt-get update -y && apt upgrade -y && apt-get install -y \ + build-essential \ + cmake \ + vim \ + wget \ + curl \ + net-tools \ + zlib1g-dev \ + lld \ + clang \ + locales \ + ccache \ + openssl \ + libssl-dev \ + pkg-config \ + ca-certificates \ + protobuf-compiler \ + && rm -rf /var/cache/apt/* \ + && rm -rf /var/lib/apt/lists/* \ + && update-ca-certificates \ + && locale-gen en_US.UTF-8 + +ENV LANG=en_US.UTF-8 +ENV LANGUAGE=en_US:en +ENV LC_ALL=en_US.UTF-8 +ENV PATH="/root/.cargo/bin:${PATH}" + +# Install dependencies +# TODO: install from pypi released memfabric +RUN pip install $MEMFABRIC_URL --no-cache-dir + +RUN pip install setuptools-rust wheel build --no-cache-dir + +# install rustup from rustup.rs +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ + && rustc --version && cargo --version && protoc --version + +# Install vLLM +RUN git clone --depth 1 https://github.com/vllm-project/vllm.git --branch $VLLM_TAG && \ + (cd vllm && VLLM_TARGET_DEVICE="empty" pip install -v . --no-cache-dir) && rm -rf vllm + +# TODO: install from pypi released triton-ascend +RUN pip install torch==$PYTORCH_VERSION torchvision==$TORCHVISION_VERSION --index-url https://download.pytorch.org/whl/cpu --no-cache-dir \ + && wget ${PTA_URL} && pip install "./torch_npu-2.6.0.post1-cp311-cp311-manylinux_2_28_aarch64.whl" --no-cache-dir \ + && python3 -m pip install --no-cache-dir attrs==24.2.0 numpy==1.26.4 scipy==1.13.1 decorator==5.1.1 psutil==6.0.0 pytest==8.3.2 pytest-xdist==3.6.1 pyyaml pybind11 \ + && pip install ${TRITON_ASCEND_URL} --no-cache-dir + +# Install SGLang +RUN git clone https://github.com/sgl-project/sglang --branch $SGLANG_TAG && \ + (cd sglang/python && pip install -v .[srt_npu] --no-cache-dir) && \ + (cd sglang/sgl-router && python -m build && pip install --force-reinstall dist/*.whl) && \ + rm -rf sglang + +# Install Deep-ep +RUN git clone --branch $SGLANG_KERNEL_NPU_TAG https://github.com/sgl-project/sgl-kernel-npu.git \ + && export LD_LIBRARY_PATH=${ASCEND_CANN_PATH}/latest/runtime/lib64/stub:$LD_LIBRARY_PATH && \ + source ${ASCEND_CANN_PATH}/set_env.sh && \ + cd sgl-kernel-npu && \ + bash build.sh \ + && pip install output/deep_ep*.whl --no-cache-dir \ + && cd .. && rm -rf sgl-kernel-npu \ + && cd "$(pip show deep-ep | awk '/^Location:/ {print $2}')" && ln -s deep_ep/deep_ep_cpp*.so + +CMD ["/bin/bash"] diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm new file mode 100644 index 000000000..2c3c9c0be --- /dev/null +++ b/docker/Dockerfile.rocm @@ -0,0 +1,213 @@ +# Usage (to build SGLang ROCm docker image): +# docker build --build-arg SGL_BRANCH=v0.5.2 --build-arg GPU_ARCH=gfx942 -t v0.5.2-rocm630-mi30x -f Dockerfile.rocm . +# docker build --build-arg SGL_BRANCH=v0.5.2 --build-arg GPU_ARCH=gfx942-rocm700 -t v0.5.2-rocm700-mi30x -f Dockerfile.rocm . +# docker build --build-arg SGL_BRANCH=v0.5.2 --build-arg GPU_ARCH=gfx950 -t v0.5.2-rocm700-mi35x -f Dockerfile.rocm . + + +# Default base images +ARG BASE_IMAGE_942="rocm/sgl-dev:vllm20250114" +ARG BASE_IMAGE_942_ROCM700="rocm/sgl-dev:rocm7-vllm-20250904" +ARG BASE_IMAGE_950="rocm/sgl-dev:rocm7-vllm-20250904" + +# This is necessary for scope purpose +ARG GPU_ARCH=gfx950 + +# =============================== +# Base image 942 with rocm630 and args +FROM $BASE_IMAGE_942 AS gfx942 +ENV BUILD_VLLM="0" +ENV BUILD_TRITON="1" +ENV BUILD_LLVM="0" +ENV BUILD_AITER_ALL="1" +ENV BUILD_MOONCAKE="1" +ENV AITER_COMMIT="v0.1.4" +ENV NO_DEPS_FLAG="" + +# =============================== +# Base image 942 and args +FROM $BASE_IMAGE_942_ROCM700 AS gfx942-rocm700 +ENV BUILD_VLLM="0" +ENV BUILD_TRITON="0" +ENV BUILD_LLVM="0" +ENV BUILD_AITER_ALL="1" +ENV BUILD_MOONCAKE="1" +ENV AITER_COMMIT="v0.1.5" +ENV NO_DEPS_FLAG="" + +# =============================== +# Base image 950 and args +FROM $BASE_IMAGE_950 AS gfx950 +ENV BUILD_VLLM="0" +ENV BUILD_TRITON="0" +ENV BUILD_LLVM="0" +ENV BUILD_AITER_ALL="1" +ENV BUILD_MOONCAKE="1" +ENV AITER_COMMIT="v0.1.5" +ENV NO_DEPS_FLAG="--no-deps" + +# =============================== +# Chosen arch and args +FROM ${GPU_ARCH} + +# This is necessary for scope purpose, again +ARG GPU_ARCH=gfx950 +ENV GPU_ARCH_LIST=${GPU_ARCH%-*} + +ARG SGL_REPO="https://github.com/sgl-project/sglang.git" +ARG SGL_DEFAULT="main" +ARG SGL_BRANCH=${SGL_DEFAULT} + +ARG TRITON_REPO="https://github.com/ROCm/triton.git" +ARG TRITON_COMMIT="improve_fa_decode_3.0.0" + +ARG AITER_REPO="https://github.com/ROCm/aiter.git" + +ARG LLVM_REPO="https://github.com/jrbyrnes/llvm-project.git" +ARG LLVM_BRANCH="MainOpSelV2" +ARG LLVM_COMMIT="6520ace8227ffe2728148d5f3b9872a870b0a560" + +ARG MOONCAKE_REPO="https://github.com/kvcache-ai/Mooncake.git" +ARG MOONCAKE_COMMIT="dcdf1c784b40aa6975a8ed89fe26321b028e40e8" + +USER root + +# Install some basic utilities +RUN python -m pip install --upgrade pip && pip install setuptools_scm +RUN apt-get purge -y sccache; python -m pip uninstall -y sccache; rm -f "$(which sccache)" + +WORKDIR /sgl-workspace + +# ----------------------- +# llvm +RUN if [ "$BUILD_LLVM" = "1" ]; then \ + ENV HIP_CLANG_PATH="/sgl-workspace/llvm-project/build/bin/" \ + git clone --single-branch ${LLVM_REPO} -b ${LLVM_BRANCH} \ + && cd llvm-project \ + && git checkout ${LLVM_COMMIT} \ + && mkdir build \ + && cd build \ + && cmake -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=1 -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" -DLLVM_ENABLE_PROJECTS="clang;lld;" -DLLVM_ENABLE_RUNTIMES="compiler-rt" ../llvm \ + && make -j$(nproc); \ + fi + +# ----------------------- + +# ----------------------- +# AITER +RUN pip uninstall -y aiter +RUN git clone ${AITER_REPO} \ + && cd aiter \ + && git checkout ${AITER_COMMIT} \ + && git submodule update --init --recursive +RUN cd aiter \ + && if [ "$BUILD_AITER_ALL" = "1" ] && [ "$BUILD_LLVM" = "1" ]; then \ + HIP_CLANG_PATH=/sgl-workspace/llvm-project/build/bin/ PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop; \ + elif [ "$BUILD_AITER_ALL" = "1" ]; then \ + PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop; \ + else \ + GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop; \ + fi + +# ----------------------- +# Triton +RUN if [ "$BUILD_TRITON" = "1" ]; then \ + pip uninstall -y triton \ + && git clone ${TRITON_REPO} \ + && cd triton \ + && git checkout ${TRITON_COMMIT} \ + && cd python \ + && python setup.py install; \ + fi + +# ----------------------- +# Build vLLM +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="9f6b92db47c3444b7a7d67451ba0c3a2d6af4c2c" +RUN if [ "$BUILD_VLLM" = "1" ]; then \ + git clone ${VLLM_REPO} \ + && cd vllm \ + && git checkout ${VLLM_BRANCH} \ + && python -m pip install -r requirements/rocm.txt \ + && python setup.py clean --all \ + && python setup.py develop; \ + fi + +# ----------------------- +# Build Mooncake +ENV PATH=$PATH:/usr/local/go/bin + +RUN if [ "$BUILD_MOONCAKE" = "1" ]; then \ + apt update && apt install -y zip unzip wget && \ + apt install -y gcc make libtool autoconf librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool libibverbs-dev rdma-core && \ + apt install -y openssh-server openmpi-bin openmpi-common libopenmpi-dev && \ + git clone ${MOONCAKE_REPO} && \ + cd Mooncake && \ + git checkout ${MOONCAKE_COMMIT} && \ + git submodule update --init --recursive && \ + bash dependencies.sh -y && \ + rm -rf /usr/local/go && \ + wget https://go.dev/dl/go1.22.2.linux-amd64.tar.gz && \ + tar -C /usr/local -xzf go1.22.2.linux-amd64.tar.gz && \ + rm go1.22.2.linux-amd64.tar.gz && \ + mkdir -p build && \ + cd build && \ + cmake .. -DUSE_ETCD=ON && \ + make -j "$(nproc)" && make install; \ + fi + + +# ----------------------- +# Build SGLang +ARG BUILD_TYPE=all + +RUN pip install IPython \ + && pip install orjson \ + && pip install python-multipart \ + && pip install torchao==0.9.0 \ + && pip install pybind11 + +RUN pip uninstall -y sgl_kernel sglang +RUN git clone ${SGL_REPO} \ + && cd sglang \ + && if [ "${SGL_BRANCH}" = ${SGL_DEFAULT} ]; then \ + echo "Using ${SGL_DEFAULT}, default branch."; \ + git checkout ${SGL_DEFAULT}; \ + else \ + echo "Using ${SGL_BRANCH} branch."; \ + git checkout ${SGL_BRANCH}; \ + fi \ + && cd sgl-kernel \ + && rm -f pyproject.toml \ + && mv pyproject_rocm.toml pyproject.toml \ + && AMDGPU_TARGET=$GPU_ARCH_LIST python setup_rocm.py install \ + && cd .. \ + && if [ "$BUILD_TYPE" = "srt" ]; then \ + python -m pip --no-cache-dir install -e "python[srt_hip]" ${NO_DEPS_FLAG}; \ + else \ + python -m pip --no-cache-dir install -e "python[all_hip]" ${NO_DEPS_FLAG}; \ + fi + +RUN python -m pip cache purge + +# Copy config files to support MI300X in virtualized environments (MI300X_VF). Symlinks will not be created in image build. +RUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \ + /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \ + -type f -name '*MI300X*' | xargs -I {} sh -c 'vf_config=$(echo "$1" | sed "s/MI300X/MI300X_VF/"); cp "$1" "$vf_config"' -- {} + +# Performance environment variable. +ENV HIP_FORCE_DEV_KERNARG=1 +ENV HSA_NO_SCRATCH_RECLAIM=1 +ENV SGLANG_SET_CPU_AFFINITY=1 +ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 +ENV NCCL_MIN_NCHANNELS=112 + +ENV SGLANG_USE_AITER=1 +ENV SGLANG_MOE_PADDING=1 +ENV VLLM_FP8_PADDING=1 +ENV VLLM_FP8_ACT_PADDING=1 +ENV VLLM_FP8_WEIGHT_PADDING=1 +ENV VLLM_FP8_REDUCE_CONV=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1 + +CMD ["/bin/bash"] diff --git a/docker/Dockerfile.router b/docker/Dockerfile.router new file mode 100644 index 000000000..ded98bb8a --- /dev/null +++ b/docker/Dockerfile.router @@ -0,0 +1,78 @@ +######################## BASE IMAGE ########################## +FROM ubuntu:24.04 AS base + +ARG PYTHON_VERSION=3.12 + +# set the environment variables +ENV PATH="/root/.local/bin:${PATH}" +ENV DEBIAN_FRONTEND=noninteractive + +# uv environment variables +ENV UV_HTTP_TIMEOUT=500 +ENV VIRTUAL_ENV="/opt/venv" +ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python +ENV UV_LINK_MODE="copy" +ENV PATH="$VIRTUAL_ENV/bin:$PATH" + + +# install dependencies +RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ + && apt update -y \ + && apt install -y curl \ + && rm -rf /var/lib/apt/lists/* \ + && apt clean + +# install uv +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +# install python +RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} + +######################### BUILD IMAGE ######################### +FROM base AS build-image + +ARG SGLANG_REPO_REF=main + +# set the environment variables +ENV PATH="/root/.cargo/bin:${PATH}" + +# install dependencies +RUN apt update -y \ + && apt install -y git build-essential libssl-dev pkg-config protobuf-compiler \ + && rm -rf /var/lib/apt/lists/* \ + && apt clean + +# install rustup from rustup.rs +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ + && rustc --version && cargo --version && protoc --version + +# pull the github repository +RUN cd /opt \ + && git clone --depth=1 https://github.com/sgl-project/sglang.git \ + && cd /opt/sglang \ + && git checkout ${SGLANG_REPO_REF} + +# working directory +WORKDIR /opt/sglang/sgl-router + +# build the rust dependencies +RUN cargo build --release \ + && uv build \ + && rm -rf /root/.cache + +######################### ROUTER IMAGE ######################### +FROM base AS router-image + +# Copy the built package from the build image +COPY --from=build-image /opt/sglang/sgl-router/dist/*.whl dist/ + +# Build the package and install +RUN uv pip install --force-reinstall dist/*.whl + +# Clean up unnecessary files to reduce the image size +RUN rm -rf /root/.cache \ + && apt purge -y --auto-remove curl + +# Set the entrypoint to the main command +ENTRYPOINT ["python3", "-m", "sglang_router.launch_router"] diff --git a/docker/Dockerfile.sagemaker b/docker/Dockerfile.sagemaker new file mode 100644 index 000000000..5fbff5096 --- /dev/null +++ b/docker/Dockerfile.sagemaker @@ -0,0 +1,6 @@ +FROM lmsysorg/sglang:latest + +COPY serve /usr/bin/serve +RUN chmod 777 /usr/bin/serve + +ENTRYPOINT [ "/usr/bin/serve" ] diff --git a/docker/Dockerfile.xeon b/docker/Dockerfile.xeon new file mode 100644 index 000000000..fdc439b30 --- /dev/null +++ b/docker/Dockerfile.xeon @@ -0,0 +1,48 @@ +FROM ubuntu:24.04 +SHELL ["/bin/bash", "-c"] + +ARG VER_SGLANG=main +ARG VER_TORCH=2.7.1 +ARG VER_TORCHVISION=0.22.1 +ARG VER_TRITON=3.3.1 + +RUN apt-get update && \ + apt-get full-upgrade -y && \ + DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \ + ca-certificates \ + git \ + curl \ + wget \ + vim \ + gcc \ + g++ \ + make + +WORKDIR /sgl-workspace + +RUN curl -fsSL -v -o miniforge.sh -O https://github.com/conda-forge/miniforge/releases/download/24.11.3-2/Miniforge3-24.11.3-2-Linux-x86_64.sh && \ + bash miniforge.sh -b -p ./miniforge3 && \ + rm -f miniforge.sh && \ + . miniforge3/bin/activate && \ + conda install -y libsqlite==3.48.0 gperftools tbb libnuma numactl + +ENV PATH=/sgl-workspace/miniforge3/bin:/sgl-workspace/miniforge3/condabin:${PATH} +ENV PIP_ROOT_USER_ACTION=ignore +ENV CONDA_PREFIX=/sgl-workspace/miniforge3 + +RUN pip config set global.index-url https://download.pytorch.org/whl/cpu && \ + pip config set global.extra-index-url https://pypi.org/simple + +RUN git clone https://github.com/sgl-project/sglang.git && \ + cd sglang && \ + git checkout ${VER_SGLANG} && \ + pip install -e "python[all_cpu]" && \ + pip install torch==${VER_TORCH} torchvision==${VER_TORCHVISION} triton==${VER_TRITON} --force-reinstall && \ + cd sgl-kernel && \ + cp pyproject_cpu.toml pyproject.toml && \ + pip install . + +ENV SGLANG_USE_CPU_ENGINE=1 +ENV LD_PRELOAD=/sgl-workspace/miniforge3/lib/libiomp5.so:/sgl-workspace/miniforge3/lib/libtcmalloc.so:/sgl-workspace/miniforge3/lib/libtbbmalloc.so.2 + +WORKDIR /sgl-workspace/sglang diff --git a/docker/compose.yaml b/docker/compose.yaml new file mode 100644 index 000000000..f7ff1fbd5 --- /dev/null +++ b/docker/compose.yaml @@ -0,0 +1,35 @@ +services: + sglang: + image: lmsysorg/sglang:latest + container_name: sglang + volumes: + - ${HOME}/.cache/huggingface:/root/.cache/huggingface + # If you use modelscope, you need mount this directory + # - ${HOME}/.cache/modelscope:/root/.cache/modelscope + restart: always + network_mode: host # required by RDMA + privileged: true # required by RDMA + # Or you can only publish port 30000 + # ports: + # - 30000:30000 + environment: + HF_TOKEN: + # if you use modelscope to download model, you need set this environment + # - SGLANG_USE_MODELSCOPE: true + entrypoint: python3 -m sglang.launch_server + command: --model-path meta-llama/Llama-3.1-8B-Instruct + --host 0.0.0.0 + --port 30000 + ulimits: + memlock: -1 + stack: 67108864 + ipc: host + healthcheck: + test: ["CMD-SHELL", "curl -f http://localhost:30000/health || exit 1"] + deploy: + resources: + reservations: + devices: + - driver: nvidia + device_ids: ["0"] + capabilities: [gpu] diff --git a/docker/k8s-sglang-distributed-sts.yaml b/docker/k8s-sglang-distributed-sts.yaml new file mode 100644 index 000000000..4252363c7 --- /dev/null +++ b/docker/k8s-sglang-distributed-sts.yaml @@ -0,0 +1,103 @@ +# Two Nodes Sglang example + +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: distributed-sglang +spec: + replicas: 2 # number of nodes/pods to run distributed sglang + selector: + matchLabels: + app: distributed-sglang + serviceName: "" + template: + metadata: + labels: + app: distributed-sglang + spec: + containers: + - name: sglang-container + image: docker.io/lmsysorg/sglang:latest + imagePullPolicy: Always # image may be replaced by official CI versioned image + command: + - /bin/bash + - -c + # please modify the sglang serving arguments below, as necessary. + # NOTE: the --expert-parallel-size is for MoE model like DeepSeek-R1 + args: + - | + python3 -m sglang.launch_server \ + --model /llm-folder \ + --dist-init-addr sglang-master-pod:5000 \ + --tensor-parallel-size 16 \ + --nnodes 2 \ + --node-rank $POD_INDEX \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 8000 \ + --enable-metrics \ + --expert-parallel-size 16 + env: + - name: POD_INDEX # reflects the node-rank + valueFrom: + fieldRef: + apiVersion: v1 + fieldPath: metadata.labels['apps.kubernetes.io/pod-index'] + - name: NCCL_DEBUG + value: INFO + resources: + limits: + nvidia.com/gpu: "8" + requests: + volumeMounts: + - mountPath: /dev/shm + name: dshm + - mountPath: /llm-folder + name: llm + securityContext: + privileged: true # to leverage RDMA/InfiniBand device, co-work with HostNetwork=true + hostNetwork: true + volumes: + - emptyDir: + medium: Memory + sizeLimit: 10Gi + name: dshm + - hostPath: + path: /llm-folder # replace with PVC or hostPath with your model weights + type: DirectoryOrCreate + name: llm + #- persistentVolumeClaim: + # claimName: llm-pvc + # name: llm +--- +apiVersion: v1 +kind: Service +metadata: + name: sglang-master-pod +spec: + type: ClusterIP + selector: + app: distributed-sglang + apps.kubernetes.io/pod-index: "0" + ports: + - name: dist-port + port: 5000 + targetPort: 5000 +--- +# the serving service +apiVersion: v1 +kind: Service +metadata: + name: sglang-serving-on-master +spec: + type: NodePort + selector: + app: distributed-sglang + apps.kubernetes.io/pod-index: "0" + ports: + - name: serving + port: 8000 + targetPort: 8000 + - name: metrics + port: 8080 + targetPort: 8080 diff --git a/docker/k8s-sglang-service.yaml b/docker/k8s-sglang-service.yaml new file mode 100644 index 000000000..866d50be9 --- /dev/null +++ b/docker/k8s-sglang-service.yaml @@ -0,0 +1,117 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: llama-31-8b-sglang +spec: + accessModes: + - ReadWriteMany + resources: + requests: + storage: 30Gi + storageClassName: default # change this to your preferred storage class + volumeMode: Filesystem +--- +apiVersion: node.k8s.io/v1 +kind: RuntimeClass +metadata: + name: nvidia +handler: nvidia +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: meta-llama-31-8b-instruct-sglang +spec: + replicas: 1 + strategy: + type: Recreate + selector: + matchLabels: + app: meta-llama-31-8b-instruct-sglang + template: + metadata: + labels: + app: meta-llama-31-8b-instruct-sglang + model: meta-llama-31-8b-instruct + engine: sglang + spec: + restartPolicy: Always + runtimeClassName: nvidia + containers: + - name: meta-llama-31-8b-instruct-sglang + image: docker.io/lmsysorg/sglang:latest + imagePullPolicy: Always # IfNotPresent or Never + ports: + - containerPort: 30000 + command: ["python3", "-m", "sglang.launch_server"] + args: + [ + "--model-path", + "meta-llama/Llama-3.1-8B-Instruct", + "--host", + "0.0.0.0", + "--port", + "30000", + ] + env: + - name: HF_TOKEN + value: + resources: + limits: + nvidia.com/gpu: 1 + cpu: 8 + memory: 40Gi + requests: + cpu: 2 + memory: 16Gi + nvidia.com/gpu: 1 + volumeMounts: + - name: shm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + - name: localtime + mountPath: /etc/localtime + readOnly: true + livenessProbe: + httpGet: + path: /health + port: 30000 + initialDelaySeconds: 120 + periodSeconds: 15 + timeoutSeconds: 10 + failureThreshold: 3 + readinessProbe: + httpGet: + path: /health_generate + port: 30000 + initialDelaySeconds: 120 + periodSeconds: 15 + timeoutSeconds: 10 + failureThreshold: 3 + successThreshold: 1 + volumes: + - name: shm + emptyDir: + medium: Memory + sizeLimit: 10Gi + - name: hf-cache + persistentVolumeClaim: + claimName: llama-31-8b-sglang + - name: localtime + hostPath: + path: /etc/localtime + type: File +--- +apiVersion: v1 +kind: Service +metadata: + name: meta-llama-31-8b-instruct-sglang +spec: + selector: + app: meta-llama-31-8b-instruct-sglang + ports: + - protocol: TCP + port: 80 # port on host + targetPort: 30000 # port in container + type: LoadBalancer # change to ClusterIP if needed diff --git a/docker/serve b/docker/serve new file mode 100755 index 000000000..493ecbd23 --- /dev/null +++ b/docker/serve @@ -0,0 +1,31 @@ +#!/bin/bash + +echo "Starting server" + +SERVER_ARGS="--host 0.0.0.0 --port 8080" + +if [ -n "$TENSOR_PARALLEL_DEGREE" ]; then + SERVER_ARGS="${SERVER_ARGS} --tp-size ${TENSOR_PARALLEL_DEGREE}" +fi + +if [ -n "$DATA_PARALLEL_DEGREE" ]; then + SERVER_ARGS="${SERVER_ARGS} --dp-size ${DATA_PARALLEL_DEGREE}" +fi + +if [ -n "$EXPERT_PARALLEL_DEGREE" ]; then + SERVER_ARGS="${SERVER_ARGS} --ep-size ${EXPERT_PARALLEL_DEGREE}" +fi + +if [ -n "$MEM_FRACTION_STATIC" ]; then + SERVER_ARGS="${SERVER_ARGS} --mem-fraction-static ${MEM_FRACTION_STATIC}" +fi + +if [ -n "$QUANTIZATION" ]; then + SERVER_ARGS="${SERVER_ARGS} --quantization ${QUANTIZATION}" +fi + +if [ -n "$CHUNKED_PREFILL_SIZE" ]; then + SERVER_ARGS="${SERVER_ARGS} --chunked-prefill-size ${CHUNKED_PREFILL_SIZE}" +fi + +python3 -m sglang.launch_server --model-path /opt/ml/model $SERVER_ARGS diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 000000000..6b8792c42 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,58 @@ +# Minimal Makefile for Sphinx documentation +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SPHINXAUTOBUILD ?= sphinx-autobuild +SOURCEDIR = . +BUILDDIR = _build +PORT ?= 8003 + +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + @echo "" + @echo "Additional targets:" + @echo " serve to build and serve documentation with auto-build and live reload" + +# Compile Notebook files and record execution time +compile: + @set -e; \ + echo "Starting Notebook compilation..."; \ + mkdir -p logs; \ + echo "Notebook execution timings:" > logs/timing.log; \ + START_TOTAL=$$(date +%s); \ + find $(SOURCEDIR) -path "*/_build/*" -prune -o -name "*.ipynb" -print0 | \ + parallel -0 -j3 --halt soon,fail=1 ' \ + NB_NAME=$$(basename {}); \ + START_TIME=$$(date +%s); \ + retry --delay=0 --times=2 -- \ + jupyter nbconvert --to notebook --execute --inplace "{}" \ + --ExecutePreprocessor.timeout=600 \ + --ExecutePreprocessor.kernel_name=python3; \ + RET_CODE=$$?; \ + END_TIME=$$(date +%s); \ + ELAPSED_TIME=$$((END_TIME - START_TIME)); \ + echo "$${NB_NAME}: $${ELAPSED_TIME}s" >> logs/timing.log; \ + exit $$RET_CODE' || exit 1; \ + END_TOTAL=$$(date +%s); \ + TOTAL_ELAPSED=$$((END_TOTAL - START_TOTAL)); \ + echo "---------------------------------" >> logs/timing.log; \ + echo "Total execution time: $${TOTAL_ELAPSED}s" >> logs/timing.log; \ + echo "All Notebook execution timings:" && cat logs/timing.log + +# Serve documentation with auto-build and live reload +serve: + @echo "Starting auto-build server at http://0.0.0.0:$(PORT)" + @$(SPHINXAUTOBUILD) "$(SOURCEDIR)" "$(BUILDDIR)/html" \ + --host 0.0.0.0 \ + --port $(PORT) \ + --watch $(SOURCEDIR) \ + --re-ignore ".*\.(ipynb_checkpoints|pyc|pyo|pyd|git)" + +.PHONY: help Makefile compile clean serve + +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +clean: + find . -name "*.ipynb" -exec nbstripout {} \; + rm -rf $(BUILDDIR) + rm -rf logs diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 000000000..2edb2c619 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,55 @@ +# SGLang Documentation + +We recommend new contributors start from writing documentation, which helps you quickly understand SGLang codebase. +Most documentation files are located under the `docs/` folder. + +## Docs Workflow + +### Install Dependency + +```bash +apt-get update && apt-get install -y pandoc parallel retry +pip install -r requirements.txt +``` + +### Update Documentation + +Update your Jupyter notebooks in the appropriate subdirectories under `docs/`. If you add new files, remember to update `index.rst` (or relevant `.rst` files) accordingly. + +- **`pre-commit run --all-files`** manually runs all configured checks, applying fixes if possible. If it fails the first time, re-run it to ensure lint errors are fully resolved. Make sure your code passes all checks **before** creating a Pull Request. + +```bash +# 1) Compile all Jupyter notebooks +make compile # This step can take a long time (10+ mins). You can consider skipping this step if you can make sure your added files are correct. +make html + +# 2) Compile and Preview documentation locally with auto-build +# This will automatically rebuild docs when files change +# Open your browser at the displayed port to view the docs +bash serve.sh + +# 2a) Alternative ways to serve documentation +# Directly use make serve +make serve +# With custom port +PORT=8080 make serve + +# 3) Clean notebook outputs +# nbstripout removes notebook outputs so your PR stays clean +pip install nbstripout +find . -name '*.ipynb' -exec nbstripout {} \; + +# 4) Pre-commit checks and create a PR +# After these checks pass, push your changes and open a PR on your branch +pre-commit run --all-files +``` +--- + +## Documentation Style Guidelines + +- For common functionalities, we prefer **Jupyter Notebooks** over Markdown so that all examples can be executed and validated by our docs CI pipeline. For complex features (e.g., distributed serving), Markdown is preferred. +- Keep in mind the documentation execution time when writing interactive Jupyter notebooks. Each interactive notebook will be run and compiled against every commit to ensure they are runnable, so it is important to apply some tips to reduce the documentation compilation time: + - Use small models (e.g., `qwen/qwen2.5-0.5b-instruct`) for most cases to reduce server launch time. + - Reuse the launched server as much as possible to reduce server launch time. +- Do not use absolute links (e.g., `https://docs.sglang.ai/get_started/install.html`). Always prefer relative links (e.g., `../get_started/install.md`). +- Follow the existing examples to learn how to launch a server, send a query and other common styles. diff --git a/docs/_static/css/custom_log.css b/docs/_static/css/custom_log.css new file mode 100644 index 000000000..61f65d019 --- /dev/null +++ b/docs/_static/css/custom_log.css @@ -0,0 +1,29 @@ +.output_area { + color: #615656; +} + +table.autosummary td { + width: 50% + } + + img.align-center { + display: block; + margin-left: auto; + margin-right: auto; +} + +.output_area.stderr { + color: #d3d3d3 !important; +} + +.output_area.stdout { + color: #d3d3d3 !important; +} + +div.output_area.stderr { + color: #d3d3d3 !important; +} + +div.output_area.stdout { + color: #d3d3d3 !important; +} diff --git a/docs/_static/css/readthedocs.css b/docs/_static/css/readthedocs.css new file mode 100644 index 000000000..aca6649b4 --- /dev/null +++ b/docs/_static/css/readthedocs.css @@ -0,0 +1,9 @@ +table.autosummary td { + width: 50% +} + +img.align-center { + display: block; + margin-left: auto; + margin-right: auto; +} diff --git a/docs/_static/image/logo.ico b/docs/_static/image/logo.ico new file mode 100644 index 000000000..7a0b5b2f0 Binary files /dev/null and b/docs/_static/image/logo.ico differ diff --git a/docs/_static/image/logo.png b/docs/_static/image/logo.png new file mode 100644 index 000000000..2a8bc258f Binary files /dev/null and b/docs/_static/image/logo.png differ diff --git a/docs/advanced_features/attention_backend.md b/docs/advanced_features/attention_backend.md new file mode 100644 index 000000000..68e4318d8 --- /dev/null +++ b/docs/advanced_features/attention_backend.md @@ -0,0 +1,104 @@ +# Attention Backend + +SGLang supports multiple attention backends. Each of them has different pros and cons. +You can test them according to your needs. + +## Supporting matrix for different attention backends + +| **Backend** | **Page Size > 1** | **Spec Decoding** | **MLA** | **Sliding Window** | **MultiModal** | +|--------------------------|-------------------|-------------------|---------|--------------------|----------------| +| **FlashInfer** | ❌ | ✅ | ✅ | ✅ | ✅ | +| **FA3** | ✅ | ✅ | ✅ | ✅ | ✅ | +| **Triton** | ❌ | ✅ | ✅ | ✅ | ❌ | +| **Torch Native** | ❌ | ❌ | ✅ | ❌ | ❌ | +| **FlashMLA** | ✅ | ✅ | ✅ | ❌ | ❌ | +| **TRTLLM MLA** | ✅ | ❌ | ✅ | ✅ | ❌ | +| **Ascend** | ✅ | ❌ | ✅ | ❌ | ❌ | +| **Wave** | ✅ | ❌ | ❌ | ❌ | ❌ | + +**Notes:** +- TRTLLM MLA only implements decode operations. For prefill operations (including multimodal inputs), it falls back to FlashInfer MLA backend. + +Note: Every kernel backend is compatible with a page size > 1 by specifying an argument such as `--page-size 16`. +This is because a page size of 16 can be converted to a page size of 1 in the kernel backend. +The "❌" and "✅" symbols in the table above under "Page Size > 1" indicate whether the kernel actually operates with a page size greater than 1, rather than treating a page size of 16 as a page size of 1. + +## User guide + +### Launch command for different attention backends. + +- FlashInfer (Default for Non-Hopper Machines, e.g., A100, A40) +```bash +python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3.1-8B-Instruct --attention-backend flashinfer +python3 -m sglang.launch_server --tp 8 --model deepseek-ai/DeepSeek-V3 --attention-backend flashinfer --trust-remote-code +``` + +- FlashAttention 3 (Default for Hopper Machines, e.g., H100, H200, H20) +```bash +python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3.1-8B-Instruct --attention-backend fa3 +python3 -m sglang.launch_server --tp 8 --model deepseek-ai/DeepSeek-V3 --trust-remote-code --attention-backend fa3 +``` + +- Triton +```bash +python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3.1-8B-Instruct --attention-backend triton +python3 -m sglang.launch_server --tp 8 --model deepseek-ai/DeepSeek-V3 --attention-backend triton --trust-remote-code +``` + +- Torch Native +```bash +python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3.1-8B-Instruct --attention-backend torch_native +``` + +- FlashMLA +```bash +python3 -m sglang.launch_server --tp 8 --model deepseek-ai/DeepSeek-R1 --attention-backend flashmla --trust-remote-code +python3 -m sglang.launch_server --tp 8 --model deepseek-ai/DeepSeek-R1 --attention-backend flashmla --kv-cache-dtype fp8_e4m3 --trust-remote-code +``` + +- TRTLLM MLA (Optimized for Blackwell Architecture, e.g., B200) +```bash +python3 -m sglang.launch_server --tp 8 --model deepseek-ai/DeepSeek-R1 --attention-backend trtllm_mla --trust-remote-code +``` + +- TRTLLM MLA with FP8 KV Cache (Higher concurrency, lower memory footprint) +```bash +python3 -m sglang.launch_server --tp 8 --model deepseek-ai/DeepSeek-R1 --attention-backend trtllm_mla --kv-cache-dtype fp8_e4m3 --trust-remote-code +``` + +- Ascend +```bash +python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3.1-8B-Instruct --attention-backend ascend +``` + +- Wave +```bash +python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3.1-8B-Instruct --attention-backend wave +``` + +## Steps to add a new attention backend +To add a new attention backend, you can learn from the existing backends +(`python/sglang/srt/layers/attention/triton_backend.py`, `python/sglang/srt/layers/attention/flashattention_backend.py`) +and follow the steps below. + +1. Run without cuda graph. Support the two forward functions + - forward_extend + - Will be used for prefill, prefill with KV cache, and target verification + - It will be called once per layer + - forward_decode + - Will be used for normal decode, and draft decode + - It will be called once per layer + - init_forward_metadata + - Initialize the class and common metadata shared by all layers + - Call the plan function for optimizations like split_kv + - It will be called once per forward +2. Run with cuda graph. It has two phases (capture and replay) and you need to implement three functions + - init_cuda_graph_state + - It will be called once during life time + - Create all common shared buffers + - init_forward_metadata_capture_cuda_graph + - It will be called before capturing a cuda graph + - It is similar to init_forward_metadata but write the medatada to some pre-defined buffers + - init_forward_metadata_replay_cuda_graph + - It will be called before replaying a cuda graph + - This function is in the critical path and needs to be fast diff --git a/docs/advanced_features/hyperparameter_tuning.md b/docs/advanced_features/hyperparameter_tuning.md new file mode 100644 index 000000000..e15ddd21c --- /dev/null +++ b/docs/advanced_features/hyperparameter_tuning.md @@ -0,0 +1,77 @@ +# Hyperparameter Tuning + +## Achieving high throughput for offline batch inference + +Achieving a large batch size is the most important thing for attaining high throughput in offline batch inference. +When the server is running at full load in a steady state, look for the following in the log: + +```Decode batch. #running-req: 233, #token: 370959, token usage: 0.82, cuda graph: True, gen throughput (token/s): 4594.01, #queue-req: 317``` + +### Adjust the request submission speed to control `#queue-req` + +`#queue-req` indicates the number of requests in the queue. +If you frequently see `#queue-req: 0`, it suggests that your client code is submitting requests too slowly. +A healthy range for `#queue-req` is `100 - 2000`. +However, avoid making `#queue-req` too large, as this will increase the scheduling overhead on the server. + +### Achieve a high `token usage` + +`token usage` indicates the KV cache memory utilization of the server. `token usage > 0.9` means good utilization. + +If you frequently see `token usage < 0.9` and `#queue-req > 0`, it means the server is too conservative about taking in new requests. You can decrease `--schedule-conservativeness` to a value like 0.3. +The case of a server being too conservative can happen when users send many requests with a large `max_new_tokens` but the requests stop very early due to EOS or stop strings. + +On the other hand, if you see `token usage` very high and you frequently see warnings like +`KV cache pool is full. Retract requests. #retracted_reqs: 1, #new_token_ratio: 0.9998 -> 1.0000`, you can increase `--schedule-conservativeness` to a value like 1.3. +If you see `KV cache pool is full. Retract requests.` occasionally but not frequently, it is okay. + +### Tune `--mem-fraction-static` to increase KV cache pool capacity +SGLang allocates memory as follows: + +Total memory usage = model weights + KV cache pool + CUDA graph buffers + activations + +The `--mem-fraction-static` parameter determines how much memory is allocated to the first two components: + +mem_fraction_static = (model weights + KV cache pool) / GPU memory capacity + +To support higher concurrency, you should maximize the KV cache pool capacity by setting `--mem-fraction-static` as high as possible while still reserving enough memory for activations and CUDA graph buffers. + +SGLang uses simple heuristics to set the default value of `--mem-fraction-static`, but you can optimize it for your use cases. +As a rule of thumb, reserving 5–8 GB of memory for activations is typically sufficient. You can check this by inspecting the logs just before the server is ready. +Look for log entries like this: + +``` +[2025-08-11 17:17:03] max_total_num_tokens=665690, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=4096, context_len=65536, available_gpu_mem=13.50 GB +``` + +Check the `available_gpu_mem` value. +- If it is between 5–8 GB, the setting is good. +- If it is too high (e.g., 10 - 20 GB), increase `--mem-fraction-static` to allocate more memory to the KV cache. +- If it is too low, you risk out-of-memory (OOM) errors later, so decrease `--mem-fraction-static`. + +Another straightforward approach is to increase `--mem-fraction-static` in increments of 0.01 until you encounter OOM errors for your workloads. + +### Avoid out-of-memory errors by tuning `--chunked-prefill-size`, `--mem-fraction-static`, and `--max-running-requests` + +If you encounter out-of-memory (OOM) errors, you can adjust the following parameters: + +- If OOM occurs during prefill, try reducing `--chunked-prefill-size` to `4096` or `2048`. This saves memory but slows down the prefill speed for long prompts. +- If OOM occurs during decoding, try lowering `--max-running-requests`. +- You can also reduce `--mem-fraction-static` to a smaller value, such as 0.8 or 0.7. This decreases the memory usage of the KV cache memory pool and helps prevent OOM errors during both prefill and decoding. However, it limits maximum concurrency and reduces peak throughput. + +### Tune `--cuda-graph-max-bs` +By default, CUDA graph is enabled only for small batch sizes (e.g., less than 160 or 256). +However, for some models, especially at large tensor parallelism sizes, CUDA graph can be useful for batch sizes up to 512 or 768. +Therefore, it may be beneficial to increase `--cuda-graph-max-bs` to a larger value. +Note that CUDA graph consumes more memory, so you may need to reduce `--mem-fraction-static` at the same time. + +### Tune `--dp-size` and `--tp-size` + +Data parallelism is better for throughput. When there is enough GPU memory, always favor data parallelism for throughput. Refer to [sglang router](../advanced_features/router.md) for a better data parallelism rather than using `dp_size` parameter. + +### Try other options + +- `torch.compile` accelerates small models on small batch sizes. You can enable it with `--enable-torch-compile`. +- Try other quantization (e.g. FP8 quantization with `--quantization fp8`) +- Try other parallelism strategies (e.g. [expert parallelism](https://lmsys.org/blog/2025-05-05-large-scale-ep/)) or DP attention for deepseek models (with `--enable-dp-attention --dp-size 8`). +- If the workload has many shared prefixes, try `--schedule-policy lpm`. Here, `lpm` stands for longest prefix match. It reorders requests to encourage more cache hits but introduces more scheduling overhead. diff --git a/docs/advanced_features/lora.ipynb b/docs/advanced_features/lora.ipynb new file mode 100644 index 000000000..1925baffc --- /dev/null +++ b/docs/advanced_features/lora.ipynb @@ -0,0 +1,529 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# LoRA Serving" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SGLang enables the use of [LoRA adapters](https://arxiv.org/abs/2106.09685) with a base model. By incorporating techniques from [S-LoRA](https://arxiv.org/pdf/2311.03285) and [Punica](https://arxiv.org/pdf/2310.18547), SGLang can efficiently support multiple LoRA adapters for different sequences within a single batch of inputs." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Arguments for LoRA Serving" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The following server arguments are relevant for multi-LoRA serving:\n", + "\n", + "* `enable_lora`: Enable LoRA support for the model. This argument is automatically set to True if `--lora-paths` is provided for backward compatibility.\n", + "\n", + "* `lora_paths`: The list of LoRA adapters to load. Each adapter must be specified in one of the following formats: | = | JSON with schema {\"lora_name\":str,\"lora_path\":str,\"pinned\":bool}.\n", + "\n", + "* `max_loras_per_batch`: Maximum number of adaptors used by each batch. This argument can affect the amount of GPU memory reserved for multi-LoRA serving, so it should be set to a smaller value when memory is scarce. Defaults to be 8.\n", + "\n", + "* `max_loaded_loras`: If specified, it limits the maximum number of LoRA adapters loaded in CPU memory at a time. The value must be greater than or equal to `max-loras-per-batch`.\n", + "\n", + "* `lora_backend`: The backend of running GEMM kernels for Lora modules. Currently we only support Triton LoRA backend. In the future, faster backend built upon Cutlass or Cuda kernels will be added.\n", + "\n", + "* `max_lora_rank`: The maximum LoRA rank that should be supported. If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. This argument is needed when you expect to dynamically load adapters of larger LoRA rank after server startup.\n", + "\n", + "* `lora_target_modules`: The union set of all target modules where LoRA should be applied (e.g., `q_proj`, `k_proj`, `gate_proj`). If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. This argument is needed when you expect to dynamically load adapters of different target modules after server startup. You can also set it to `all` to enable LoRA for all supported modules. However, enabling LoRA on additional modules introduces a minor performance overhead. If your application is performance-sensitive, we recommend only specifying the modules for which you plan to load adapters.\n", + "\n", + "* `tp_size`: LoRA serving along with Tensor Parallelism is supported by SGLang. `tp_size` controls the number of GPUs for tensor parallelism. More details on the tensor sharding strategy can be found in [S-Lora](https://arxiv.org/pdf/2311.03285) paper.\n", + "\n", + "From client side, the user needs to provide a list of strings as input batch, and a list of adaptor names that each input sequence corresponds to." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Usage\n", + "\n", + "### Serving Single Adaptor" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import requests\n", + "\n", + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, terminate_process" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", + " --enable-lora \\\n", + " --lora-paths lora0=algoprog/fact-generation-llama-3.1-8b-instruct-lora \\\n", + " --max-loras-per-batch 1 --lora-backend triton \\\n", + " --log-level warning \\\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://127.0.0.1:{port}\"\n", + "json_data = {\n", + " \"text\": [\n", + " \"List 3 countries and their capitals.\",\n", + " \"List 3 countries and their capitals.\",\n", + " ],\n", + " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", + " # The first input uses lora0, and the second input uses the base model\n", + " \"lora_path\": [\"lora0\", None],\n", + "}\n", + "response = requests.post(\n", + " url + \"/generate\",\n", + " json=json_data,\n", + ")\n", + "print(f\"Output 0: {response.json()[0]['text']}\")\n", + "print(f\"Output 1: {response.json()[1]['text']}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Serving Multiple Adaptors" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", + " --enable-lora \\\n", + " --lora-paths lora0=algoprog/fact-generation-llama-3.1-8b-instruct-lora \\\n", + " lora1=Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16 \\\n", + " --max-loras-per-batch 2 --lora-backend triton \\\n", + " --log-level warning \\\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://127.0.0.1:{port}\"\n", + "json_data = {\n", + " \"text\": [\n", + " \"List 3 countries and their capitals.\",\n", + " \"List 3 countries and their capitals.\",\n", + " ],\n", + " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", + " # The first input uses lora0, and the second input uses lora1\n", + " \"lora_path\": [\"lora0\", \"lora1\"],\n", + "}\n", + "response = requests.post(\n", + " url + \"/generate\",\n", + " json=json_data,\n", + ")\n", + "print(f\"Output 0: {response.json()[0]['text']}\")\n", + "print(f\"Output 1: {response.json()[1]['text']}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Dynamic LoRA loading" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Instead of specifying all adapters during server startup via `--lora-paths`. You can also load & unload LoRA adapters dynamically via the `/load_lora_adapter` and `/unload_lora_adapter` API.\n", + "\n", + "When using dynamic LoRA loading, it's recommended to explicitly specify both `--max-lora-rank` and `--lora-target-modules` at startup. For backward compatibility, SGLang will infer these values from `--lora-paths` if they are not explicitly provided. However, in that case, you would have to ensure that all dynamically loaded adapters share the same shape (rank and target modules) as those in the initial `--lora-paths` or are strictly \"smaller\"." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "lora0 = \"Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16\" # rank - 4, target modules - q_proj, k_proj, v_proj, o_proj, gate_proj\n", + "lora1 = \"algoprog/fact-generation-llama-3.1-8b-instruct-lora\" # rank - 64, target modules - q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj\n", + "lora0_new = \"philschmid/code-llama-3-1-8b-text-to-sql-lora\" # rank - 256, target modules - q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj\n", + "\n", + "\n", + "# The `--target-lora-modules` param below is technically not needed, as the server will infer it from lora0 which already has all the target modules specified.\n", + "# We are adding it here just to demonstrate usage.\n", + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + " python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", + " --enable-lora \\\n", + " --cuda-graph-max-bs 2 \\\n", + " --max-loras-per-batch 2 --lora-backend triton \\\n", + " --max-lora-rank 256\n", + " --lora-target-modules all\n", + " --log-level warning\n", + " \"\"\"\n", + ")\n", + "\n", + "url = f\"http://127.0.0.1:{port}\"\n", + "wait_for_server(url)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load adapter lora0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = requests.post(\n", + " url + \"/load_lora_adapter\",\n", + " json={\n", + " \"lora_name\": \"lora0\",\n", + " \"lora_path\": lora0,\n", + " },\n", + ")\n", + "\n", + "if response.status_code == 200:\n", + " print(\"LoRA adapter loaded successfully.\", response.json())\n", + "else:\n", + " print(\"Failed to load LoRA adapter.\", response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load adapter lora1:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = requests.post(\n", + " url + \"/load_lora_adapter\",\n", + " json={\n", + " \"lora_name\": \"lora1\",\n", + " \"lora_path\": lora1,\n", + " },\n", + ")\n", + "\n", + "if response.status_code == 200:\n", + " print(\"LoRA adapter loaded successfully.\", response.json())\n", + "else:\n", + " print(\"Failed to load LoRA adapter.\", response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Check inference output:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://127.0.0.1:{port}\"\n", + "json_data = {\n", + " \"text\": [\n", + " \"List 3 countries and their capitals.\",\n", + " \"List 3 countries and their capitals.\",\n", + " ],\n", + " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", + " # The first input uses lora0, and the second input uses lora1\n", + " \"lora_path\": [\"lora0\", \"lora1\"],\n", + "}\n", + "response = requests.post(\n", + " url + \"/generate\",\n", + " json=json_data,\n", + ")\n", + "print(f\"Output from lora0: \\n{response.json()[0]['text']}\\n\")\n", + "print(f\"Output from lora1 (updated): \\n{response.json()[1]['text']}\\n\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Unload lora0 and replace it with a different adapter:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = requests.post(\n", + " url + \"/unload_lora_adapter\",\n", + " json={\n", + " \"lora_name\": \"lora0\",\n", + " },\n", + ")\n", + "\n", + "response = requests.post(\n", + " url + \"/load_lora_adapter\",\n", + " json={\n", + " \"lora_name\": \"lora0\",\n", + " \"lora_path\": lora0_new,\n", + " },\n", + ")\n", + "\n", + "if response.status_code == 200:\n", + " print(\"LoRA adapter loaded successfully.\", response.json())\n", + "else:\n", + " print(\"Failed to load LoRA adapter.\", response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Check output again:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://127.0.0.1:{port}\"\n", + "json_data = {\n", + " \"text\": [\n", + " \"List 3 countries and their capitals.\",\n", + " \"List 3 countries and their capitals.\",\n", + " ],\n", + " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", + " # The first input uses lora0, and the second input uses lora1\n", + " \"lora_path\": [\"lora0\", \"lora1\"],\n", + "}\n", + "response = requests.post(\n", + " url + \"/generate\",\n", + " json=json_data,\n", + ")\n", + "print(f\"Output from lora0: \\n{response.json()[0]['text']}\\n\")\n", + "print(f\"Output from lora1 (updated): \\n{response.json()[1]['text']}\\n\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### LoRA GPU Pinning" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Another advanced option is to specify adapters as `pinned` during loading. When an adapter is pinned, it is permanently assigned to one of the available GPU pool slots (as configured by `--max-loras-per-batch`) and will not be evicted from GPU memory during runtime. Instead, it remains resident until it is explicitly unloaded.\n", + "\n", + "This can improve performance in scenarios where the same adapter is frequently used across requests, by avoiding repeated memory transfers and reinitialization overhead. However, since GPU pool slots are limited, pinning adapters reduces the flexibility of the system to dynamically load other adapters on demand. If too many adapters are pinned, it may lead to degraded performance, or in the most extreme case (`Number of pinned adapters == max-loras-per-batch`), halt all unpinned requests. Therefore, currently SGLang limits maximal number of pinned adapters to `max-loras-per-batch - 1` to prevent unexpected starvations. \n", + "\n", + "In the example below, we start a server with `lora1` loaded as pinned, `lora2` and `lora3` loaded as regular (unpinned) adapters. Please note that, we intentionally specify `lora2` and `lora3` in two different formats to demonstrate that both are supported." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + " python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", + " --enable-lora \\\n", + " --cuda-graph-max-bs 8 \\\n", + " --max-loras-per-batch 3 --lora-backend triton \\\n", + " --max-lora-rank 256 \\\n", + " --lora-target-modules all \\\n", + " --lora-paths \\\n", + " {\"lora_name\":\"lora0\",\"lora_path\":\"Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16\",\"pinned\":true} \\\n", + " {\"lora_name\":\"lora1\",\"lora_path\":\"algoprog/fact-generation-llama-3.1-8b-instruct-lora\"} \\\n", + " lora2=philschmid/code-llama-3-1-8b-text-to-sql-lora\n", + " --log-level warning\n", + " \"\"\"\n", + ")\n", + "\n", + "\n", + "url = f\"http://127.0.0.1:{port}\"\n", + "wait_for_server(url)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can also specify adapter as pinned during dynamic adapter loading. In the example below, we reload `lora2` as pinned adapter:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = requests.post(\n", + " url + \"/unload_lora_adapter\",\n", + " json={\n", + " \"lora_name\": \"lora1\",\n", + " },\n", + ")\n", + "\n", + "response = requests.post(\n", + " url + \"/load_lora_adapter\",\n", + " json={\n", + " \"lora_name\": \"lora1\",\n", + " \"lora_path\": \"algoprog/fact-generation-llama-3.1-8b-instruct-lora\",\n", + " \"pinned\": True, # Pin the adapter to GPU\n", + " },\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Verify that the results are expected:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://127.0.0.1:{port}\"\n", + "json_data = {\n", + " \"text\": [\n", + " \"List 3 countries and their capitals.\",\n", + " \"List 3 countries and their capitals.\",\n", + " \"List 3 countries and their capitals.\",\n", + " ],\n", + " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", + " # The first input uses lora0, and the second input uses lora1\n", + " \"lora_path\": [\"lora0\", \"lora1\", \"lora2\"],\n", + "}\n", + "response = requests.post(\n", + " url + \"/generate\",\n", + " json=json_data,\n", + ")\n", + "print(f\"Output from lora0 (pinned): \\n{response.json()[0]['text']}\\n\")\n", + "print(f\"Output from lora1 (pinned): \\n{response.json()[1]['text']}\\n\")\n", + "print(f\"Output from lora2 (not pinned): \\n{response.json()[2]['text']}\\n\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Future Works\n", + "\n", + "The development roadmap for LoRA-related features can be found in this [issue](https://github.com/sgl-project/sglang/issues/2929). Other features, including Embedding Layer, Unified Paging, Cutlass backend are still under development." + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/advanced_features/observability.md b/docs/advanced_features/observability.md new file mode 100644 index 000000000..f03fb3772 --- /dev/null +++ b/docs/advanced_features/observability.md @@ -0,0 +1,35 @@ +# Observability + +## Production Metrics +SGLang exposes the following metrics via Prometheus. You can enable them by adding `--enable-metrics` when launching the server. +You can query them by: +``` +curl http://localhost:30000/metrics +``` + +See [Production Metrics](../references/production_metrics.md) for more details. + +## Logging + +By default, SGLang does not log any request contents. You can log them by using `--log-requests`. +You can control the verbosity by using `--log-request-level`. +See [Logging](server_arguments.md#logging) for more details. + +## Request Dump and Replay + +You can dump all requests and replay them later for benchmarking or other purposes. + +To start dumping, use the following command to send a request to a server: +``` +python3 -m sglang.srt.managers.configure_logging --url http://localhost:30000 --dump-requests-folder /tmp/sglang_request_dump --dump-requests-threshold 100 +``` +The server will dump the requests into a pickle file for every 100 requests. + +To replay the request dump, use `scripts/playground/replay_request_dump.py`. + +## Crash Dump and Replay +Sometimes the server might crash, and you may want to debug the cause of the crash. +SGLang supports crash dumping, which will dump all requests from the 5 minutes before the crash, allowing you to replay the requests and debug the reason later. + +To enable crash dumping, use `--crash-dump-folder /tmp/crash_dump`. +To replay the crash dump, use `scripts/playground/replay_request_dump.py`. diff --git a/docs/advanced_features/pd_disaggregation.md b/docs/advanced_features/pd_disaggregation.md new file mode 100644 index 000000000..85a5db07e --- /dev/null +++ b/docs/advanced_features/pd_disaggregation.md @@ -0,0 +1,150 @@ +# PD Disaggregation + +## Why and What is PD Disaggregation? + +Large Language Model (LLM) inference comprises two distinct phases: **Prefill** and **Decode**. The Prefill phase is computation-intensive, processing the entire input sequence, while the Decode phase is memory-intensive, managing the Key-Value (KV) cache for token generation. Traditionally, these phases are handled within a unified engine, where combined scheduling of prefill and decode batches introduces inefficiencies. To address these challenges, we introduce **Prefill and Decoding (PD) Disaggregation** in SGLang. + +### Issues with Unified Scheduling + +The conventional unified engine, which processes prefill and decode batches together, results in two significant problems: + +1. **Prefill Interruption**: Incoming prefill batches frequently interrupt ongoing decode batches, causing substantial delays in token generation. +2. **DP Attention Imbalance**: In data-parallel (DP) attention, one DP worker may process a prefill batch while another handles a decode batch simultaneously, leading to increased decode latency. + +PD Disaggregation resolves these by separating the two stages, enabling tailored optimizations for each. + +For the design details, please refer to [link](https://docs.google.com/document/d/1rQXJwKd5b9b1aOzLh98mnyMhBMhlxXA5ATZTHoQrwvc/edit?tab=t.0). + +Currently, we support Mooncake and NIXL as the transfer engine. + +## Router Integration + +For deploying PD disaggregation at scale with load balancing and fault tolerance, SGLang provides a router. The router can distribute requests between prefill and decode instances using various routing policies. For detailed information on setting up routing with PD disaggregation, including configuration options and deployment patterns, see the [SGLang Router documentation](router.md#mode-3-prefill-decode-disaggregation). + + +## Mooncake +### Requirements + +```bash +uv pip install mooncake-transfer-engine +``` + +### Usage + +### Llama Single Node + +```bash +$ python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode prefill --disaggregation-ib-device mlx5_roce0 +$ python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode decode --port 30001 --base-gpu-id 1 --disaggregation-ib-device mlx5_roce0 +$ python -m sglang_router.launch_router --pd-disaggregation --prefill http://127.0.0.1:30000 --decode http://127.0.0.1:30001 --host 0.0.0.0 --port 8000 +``` + +### DeepSeek Multi-Node + +```bash +# prefill 0 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-ib-device ${device_name} --disaggregation-mode prefill --host ${local_ip} --port 30000 --trust-remote-code --dist-init-addr ${prefill_master_ip}:5000 --nnodes 2 --node-rank 0 --tp-size 16 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep --mem-fraction-static 0.8 +# prefill 1 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-ib-device ${device_name} --disaggregation-mode prefill --host ${local_ip} --port 30000 --trust-remote-code --dist-init-addr ${prefill_master_ip}:5000 --nnodes 2 --node-rank 1 --tp-size 16 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep --mem-fraction-static 0.8 +# decode 0 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-ib-device ${device_name} --disaggregation-mode decode --host ${local_ip} --port 30001 --trust-remote-code --dist-init-addr ${decode_master_ip}:5000 --nnodes 2 --node-rank 0 --tp-size 16 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep --mem-fraction-static 0.8 --max-running-requests 128 +# decode 1 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-ib-device ${device_name} --disaggregation-mode decode --host ${local_ip} --port 30001 --trust-remote-code --dist-init-addr ${decode_master_ip}:5000 --nnodes 2 --node-rank 1 --tp-size 16 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep --mem-fraction-static 0.8 --max-running-requests 128 +``` +### Advanced Configuration + +PD Disaggregation with Mooncake supports the following environment variables for fine-grained control over system behavior. + +#### Prefill Server Configuration +| Variable | Description | Default | +|:--------:|:-----------:|:--------: +| **`SGLANG_DISAGGREGATION_THREAD_POOL_SIZE`** | Controls the total number of worker threads for KVCache transfer operations per TP rank | A dynamic value calculated by `int(0.75 * os.cpu_count()) // 8)`, which is limited to be larger than 4 and less than 12 to ensure efficiency and prevent thread race conditions | +| **`SGLANG_DISAGGREGATION_QUEUE_SIZE`** | Sets the number of parallel transfer queues. KVCache transfer requests from multiple decode instances will be sharded into these queues so that they can share the threads and the transfer bandwidth at the same time. If it is set to `1`, then we transfer requests one by one according to fcfs strategy | `4` | +| **`SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT`** | Timeout (seconds) for receiving destination KV indices during request initialization | `300` | + +If a greater mean TTFT is acceptable, you can `export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600` (10 minutes) to relax the timeout condition. +Please be aware that this setting will cause prefill instances to take a longer time to clean up the affected memory resources when a running decode node loses connection. + +#### Decode Server Configuration +| Variable | Description | Default | +|:--------:|:-----------:|:--------: +| **`SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL`** | Interval (seconds) between health checks to prefill bootstrap servers | `5.0` | +| **`SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE`** | Consecutive heartbeat failures before marking prefill server offline | `2` | +| **`SGLANG_DISAGGREGATION_WAITING_TIMEOUT`** | Timeout (seconds) for receiving KV Cache after request initialization | `300` | + +If a greater mean TTFT is acceptable, you can `export SGLANG_DISAGGREGATION_WAITING_TIMEOUT=600` (10 minutes) to relax the timeout condition. + + +## NIXL +### Requirements + +Install via pip. + +```bash +pip install nixl +``` + +Or build from source - may be required if you already have UCX installed. + +```bash +git clone https://github.com/ai-dynamo/nixl.git +cd nixl +pip install . --config-settings=setup-args="-Ducx_path=/path/to/ucx" +``` + + +### Usage + +### Llama Single Node + +```bash +$ python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode prefill --disaggregation-transfer-backend nixl +$ python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode decode --port 30001 --base-gpu-id 1 --disaggregation-transfer-backend nixl +$ python -m sglang_router.launch_router --pd-disaggregation --prefill http://127.0.0.1:30000 --decode http://127.0.0.1:30001 --host 0.0.0.0 --port 8000 +``` + +### DeepSeek Multi-Node + +```bash +# prefill 0 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-transfer-backend nixl --disaggregation-mode prefill --host ${local_ip} --port 30000 --trust-remote-code --dist-init-addr ${prefill_master_ip}:5000 --nnodes 2 --node-rank 0 --tp-size 16 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep --mem-fraction-static 0.8 +# prefill 1 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-transfer-backend nixl --disaggregation-mode prefill --host ${local_ip} --port 30000 --trust-remote-code --dist-init-addr ${prefill_master_ip}:5000 --nnodes 2 --node-rank 1 --tp-size 16 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep --mem-fraction-static 0.8 +# decode 0 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-transfer-backend nixl --disaggregation-mode decode --host ${local_ip} --port 30001 --trust-remote-code --dist-init-addr ${decode_master_ip}:5000 --nnodes 2 --node-rank 0 --tp-size 16 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep --mem-fraction-static 0.8 --max-running-requests 128 +# decode 1 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-transfer-backend nixl --disaggregation-mode decode --host ${local_ip} --port 30001 --trust-remote-code --dist-init-addr ${decode_master_ip}:5000 --nnodes 2 --node-rank 1 --tp-size 16 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep --mem-fraction-static 0.8 --max-running-requests 128 +``` + +## ASCEND + +### Usage + +Use ascend backend with [mf_adapter(download link)](https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com:443/sglang/mf_adapter-1.0.0-cp311-cp311-linux_aarch64.whl?AccessKeyId=HPUAXT4YM0U8JNTERLST&Expires=1783151861&Signature=3j10QDUjqk70enaq8lostYV2bEA%3D) and ASCEND_MF_STORE_URL being set + +```bash +pip install mf_adapter-1.0.0-cp311-cp311-linux_aarch64.whl --force-reinstall +export ASCEND_MF_STORE_URL="tcp://xxx.xx.xxx.xxx:xxxx" +``` +Use mooncake backend, more details can be found in mooncake section. +```bash +export ENABLE_ASCEND_TRANSFER_WITH_MOONCAKE=true +``` + + +### Llama Single Node + +```bash +$ python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode prefill --disaggregation-transfer-backend ascend +$ python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode decode --port 30001 --base-gpu-id 1 --disaggregation-transfer-backend ascend +$ python -m sglang_router.launch_router --pd-disaggregation --prefill http://127.0.0.1:30000 --decode http://127.0.0.1:30001 --host 0.0.0.0 --port 8000 +``` + +### DeepSeek Multi-Node + +```bash +# prefill 0 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-transfer-backend ascend --disaggregation-mode prefill --host ${local_ip} --port 30000 --trust-remote-code --dist-init-addr ${prefill_master_ip}:5000 --nnodes 1 --node-rank 0 --tp-size 16 +# decode 0 +$ python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --disaggregation-transfer-backend ascend --disaggregation-mode decode --host ${local_ip} --port 30001 --trust-remote-code --dist-init-addr ${decode_master_ip}:5000 --nnodes 1 --node-rank 0 --tp-size 16 +``` diff --git a/docs/advanced_features/quantization.md b/docs/advanced_features/quantization.md new file mode 100644 index 000000000..3a229f83d --- /dev/null +++ b/docs/advanced_features/quantization.md @@ -0,0 +1,152 @@ +# Quantization + +SGLang supports various quantization methods, including offline quantization and online dynamic quantization. + +Offline quantization loads pre-quantized model weights directly during inference. This is required for quantization methods +such as GPTQ and AWQ, which collect and pre-compute various statistics from the original weights using the calibration dataset. + +Online quantization dynamically computes scaling parameters—such as the maximum/minimum values of model weights—during runtime. +Like NVIDIA FP8 training's [delayed scaling](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#Mixed-precision-training-with-FP8) mechanism, online quantization calculates the appropriate scaling factors +on-the-fly to convert high-precision weights into a lower-precision format. + +**Note: For better performance, usability and convenience, offline quantization is recommended over online quantization.** + +If you use a pre-quantized model, do not add `--quantization` to enable online quantization at the same time. +For popular pre-quantized models, please visit [ModelCloud](https://huggingface.co/collections/ModelCloud/vortex-673743382af0a52b2a8b9fe2) +or [NeuralMagic](https://huggingface.co/collections/neuralmagic) collections on HF for some +popular quality validated quantized models. Quantized models must be validated via benchmarks post-quantization +to guard against abnormal quantization loss regressions. + +## Offline Quantization + +To load already quantized models, simply load the model weights and config. **Again, if the model has been quantized offline, +there's no need to add `--quantization` argument when starting the engine. The quantization method will be parsed from the +downloaded Hugging Face config. For example, DeepSeek V3/R1 models are already in FP8, so do not add redundant parameters.** + +```bash +python3 -m sglang.launch_server \ + --model-path hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4 \ + --port 30000 --host 0.0.0.0 +``` + +Take note, if your model is **per-channel quantized (INT8 or FP8) with per-token dynamic quantization activation**, you can opt to include `--quantization w8a8_int8` or `--quantization w8a8_fp8` to invoke the corresponding CUTLASS int8_kernel or fp8_kernel in sgl-kernel. This action will ignore the Hugging Face config's quantization settings. For instance, with `neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic`, if you execute with `--quantization w8a8_fp8`, the system will use the `W8A8Fp8Config` from SGLang to invoke the sgl-kernel, rather than the `CompressedTensorsConfig` for vLLM kernels. + +```bash +python3 -m sglang.launch_server \ + --model-path neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic \ + --quantization w8a8_fp8 \ + --port 30000 --host 0.0.0.0 +``` + +### Examples of Offline Model Quantization + +#### Using [GPTQModel](https://github.com/ModelCloud/GPTQModel) + +```bash +# install +pip install gptqmodel --no-build-isolation -v +``` + +```py +from datasets import load_dataset +from gptqmodel import GPTQModel, QuantizeConfig + +model_id = "meta-llama/Llama-3.2-1B-Instruct" +quant_path = "Llama-3.2-1B-Instruct-gptqmodel-4bit" + +calibration_dataset = load_dataset( + "allenai/c4", data_files="en/c4-train.00001-of-01024.json.gz", + split="train" + ).select(range(1024))["text"] + +quant_config = QuantizeConfig(bits=4, group_size=128) # quantization config +model = GPTQModel.load(model_id, quant_config) # load model + +model.quantize(calibration_dataset, batch_size=2) # quantize +model.save(quant_path) # save model +``` + +#### Using [LLM Compressor](https://github.com/vllm-project/llm-compressor/) + +```bash +# install +pip install llmcompressor +``` + +Here, we take quantize `meta-llama/Meta-Llama-3-8B-Instruct` to `FP8` as an example to elaborate on how to do offline quantization. + +```python +from transformers import AutoTokenizer +from llmcompressor.transformers import SparseAutoModelForCausalLM +from llmcompressor.transformers import oneshot +from llmcompressor.modifiers.quantization import QuantizationModifier + +# Step 1: Load the original model. +MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct" + +model = SparseAutoModelForCausalLM.from_pretrained( + MODEL_ID, device_map="auto", torch_dtype="auto") +tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) + +# Step 2: Perform offline quantization. +# Step 2.1: Configure the simple PTQ quantization. +recipe = QuantizationModifier( + targets="Linear", scheme="FP8_DYNAMIC", ignore=["lm_head"]) + +# Step 2.2: Apply the quantization algorithm. +oneshot(model=model, recipe=recipe) + +# Step 3: Save the model. +SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic" +model.save_pretrained(SAVE_DIR) +tokenizer.save_pretrained(SAVE_DIR) +``` + +Then, you can directly use the quantized model with `SGLang`, by using the following command: + +```bash +python3 -m sglang.launch_server \ + --model-path $PWD/Meta-Llama-3-8B-Instruct-FP8-Dynamic \ + --port 30000 --host 0.0.0.0 +``` + +## Online Quantization + +To enable online quantization, you can simply specify `--quantization` in the command line. For example, you can launch the server with the following command to enable `FP8` quantization for model `meta-llama/Meta-Llama-3.1-8B-Instruct`: + +```bash +python3 -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --quantization fp8 \ + --port 30000 --host 0.0.0.0 +``` + +Our team is working on supporting more online quantization methods. SGLang will soon support methods including but not limited to `["awq", "gptq", "marlin", "gptq_marlin", "awq_marlin", "bitsandbytes", "gguf"]`. + +SGLang also supports quantization methods based on [torchao](https://github.com/pytorch/ao). You can simply specify `--torchao-config` in the command line to support this feature. For example, if you want to enable `int4wo-128` for model `meta-llama/Meta-Llama-3.1-8B-Instruct`, you can launch the server with the following command: + +```bash +python3 -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --torchao-config int4wo-128 \ + --port 30000 --host 0.0.0.0 +``` + +SGLang supports the following quantization methods based on torchao `["int8dq", "int8wo", "fp8wo", "fp8dq-per_tensor", "fp8dq-per_row", "int4wo-32", "int4wo-64", "int4wo-128", "int4wo-256"]`. + +Note: According to [this issue](https://github.com/sgl-project/sglang/issues/2219#issuecomment-2561890230), `"int8dq"` method currently has some bugs when using together with cuda graph capture. So we suggest to disable cuda graph capture when using `"int8dq"` method. Namely, please use the following command: + +```bash +python3 -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --torchao-config int8dq \ + --disable-cuda-graph \ + --port 30000 --host 0.0.0.0 +``` + +## Reference + +- [GPTQModel](https://github.com/ModelCloud/GPTQModel) +- [LLM Compressor](https://github.com/vllm-project/llm-compressor/) +- [Torchao: PyTorch Architecture Optimization](https://github.com/pytorch/ao) +- [vLLM Quantization](https://docs.vllm.ai/en/latest/quantization/) diff --git a/docs/advanced_features/router.md b/docs/advanced_features/router.md new file mode 100644 index 000000000..4aba99f37 --- /dev/null +++ b/docs/advanced_features/router.md @@ -0,0 +1,445 @@ +# SGLang Router + +The SGLang Router is a high-performance request distribution system that routes inference requests across multiple SGLang runtime instances. It features cache-aware load balancing, fault tolerance, and support for advanced deployment patterns including data parallelism and prefill-decode disaggregation. + +## Key Features + +- **Cache-Aware Load Balancing**: Optimizes cache utilization while maintaining balanced load distribution +- **Multiple Routing Policies**: Choose from random, round-robin, cache-aware, or power-of-two policies +- **Fault Tolerance**: Automatic retry and circuit breaker mechanisms for resilient operation +- **Dynamic Scaling**: Add or remove workers at runtime without service interruption +- **Kubernetes Integration**: Native service discovery and pod management +- **Prefill-Decode Disaggregation**: Support for disaggregated serving load balancing +- **Prometheus Metrics**: Built-in observability and monitoring + +## Installation + +```bash +pip install sglang-router +``` + +## Quick Start + +To see all available options: + +```bash +python -m sglang_router.launch_server --help # Co-launch router and workers +python -m sglang_router.launch_router --help # Launch router only +``` + +## Deployment Modes + +The router supports three primary deployment patterns: + +1. **Co-launch Mode**: Router and workers launch together (simplest for single-node deployments) +2. **Separate Launch Mode**: Router and workers launch independently (best for multi-node setups) +3. **Prefill-Decode Disaggregation**: Specialized mode for disaggregated serving + +### Mode 1: Co-launch Router and Workers + +This mode launches both the router and multiple worker instances in a single command. It's the simplest deployment option and replaces the `--dp-size` argument of SGLang Runtime. + +```bash +# Launch router with 4 workers +python -m sglang_router.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --dp-size 4 \ + --host 0.0.0.0 \ + --port 30000 +``` + +#### Sending Requests + +Once the server is ready, send requests to the router endpoint: + +```python +import requests + +# Using the /generate endpoint +url = "http://localhost:30000/generate" +data = { + "text": "What is the capital of France?", + "sampling_params": { + "temperature": 0.7, + "max_new_tokens": 100 + } +} + +response = requests.post(url, json=data) +print(response.json()) + +# OpenAI-compatible endpoint +url = "http://localhost:30000/v1/chat/completions" +data = { + "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", + "messages": [{"role": "user", "content": "What is the capital of France?"}] +} + +response = requests.post(url, json=data) +print(response.json()) +``` + +### Mode 2: Separate Launch Mode + +This mode is ideal for multi-node deployments where workers run on different machines. + +#### Step 1: Launch Workers + +On each worker node: + +```bash +# Worker node 1 +python -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --host 0.0.0.0 \ + --port 8000 + +# Worker node 2 +python -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --host 0.0.0.0 \ + --port 8001 +``` + +#### Step 2: Launch Router + +On the router node: + +```bash +python -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 http://worker2:8001 \ + --host 0.0.0.0 \ + --port 30000 \ + --policy cache_aware # or random, round_robin, power_of_two +``` + +### Mode 3: Prefill-Decode Disaggregation + +This advanced mode separates prefill and decode operations for optimized performance: + +```bash +python -m sglang_router.launch_router \ + --pd-disaggregation \ + --prefill http://prefill1:8000 9000 \ + --prefill http://prefill2:8001 9001 \ + --decode http://decode1:8002 \ + --decode http://decode2:8003 \ + --prefill-policy cache_aware \ + --decode-policy round_robin +``` + +#### Understanding --prefill Arguments + +The `--prefill` flag accepts URLs with optional bootstrap ports: +- `--prefill http://server:8000` - No bootstrap port +- `--prefill http://server:8000 9000` - Bootstrap port 9000 +- `--prefill http://server:8000 none` - Explicitly no bootstrap port + +#### Policy Inheritance in PD Mode + +The router intelligently handles policy configuration for prefill and decode nodes: + +1. **Only `--policy` specified**: Both prefill and decode nodes use this policy +2. **`--policy` and `--prefill-policy` specified**: Prefill nodes use `--prefill-policy`, decode nodes use `--policy` +3. **`--policy` and `--decode-policy` specified**: Prefill nodes use `--policy`, decode nodes use `--decode-policy` +4. **All three specified**: Prefill nodes use `--prefill-policy`, decode nodes use `--decode-policy` (main `--policy` is ignored) + +Example with mixed policies: +```bash +python -m sglang_router.launch_router \ + --pd-disaggregation \ + --prefill http://prefill1:8000 + --prefill http://prefill2:8000 \ + --decode http://decode1:8001 + --decode http://decode2:8001 \ + --policy round_robin \ + --prefill-policy cache_aware # Prefill uses cache_aware and decode uses round_robin from --policy +``` + +#### PD Mode with Service Discovery + +For Kubernetes deployments with separate prefill and decode server pools: + +```bash +python -m sglang_router.launch_router \ + --pd-disaggregation \ + --service-discovery \ + --prefill-selector app=prefill-server tier=gpu \ + --decode-selector app=decode-server tier=cpu \ + --service-discovery-namespace production \ + --prefill-policy cache_aware \ + --decode-policy round_robin +``` + +## Dynamic Scaling + +The router supports runtime scaling through REST APIs: + +### Adding Workers + +```bash +# Launch a new worker +python -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --port 30001 + +# Add it to the router +curl -X POST "http://localhost:30000/add_worker?url=http://127.0.0.1:30001" +``` + +### Removing Workers + +```bash +curl -X POST "http://localhost:30000/remove_worker?url=http://127.0.0.1:30001" +``` + +**Note**: When using cache-aware routing, removed workers are cleanly evicted from the routing tree and request queues. + +## Fault Tolerance + +The router includes comprehensive fault tolerance mechanisms: + +### Retry Configuration + +```bash +python -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 http://worker2:8001 \ + --retry-max-retries 3 \ + --retry-initial-backoff-ms 100 \ + --retry-max-backoff-ms 10000 \ + --retry-backoff-multiplier 2.0 \ + --retry-jitter-factor 0.1 +``` + +### Circuit Breaker + +Protects against cascading failures: + +```bash +python -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 http://worker2:8001 \ + --cb-failure-threshold 5 \ + --cb-success-threshold 2 \ + --cb-timeout-duration-secs 30 \ + --cb-window-duration-secs 60 +``` + +**Behavior**: +- Worker is marked unhealthy after `cb-failure-threshold` consecutive failures +- Returns to service after `cb-success-threshold` successful health checks +- Circuit breaker can be disabled with `--disable-circuit-breaker` + +## Routing Policies + +The router supports multiple routing strategies: + +### 1. Random Routing +Distributes requests randomly across workers. + +```bash +--policy random +``` + +### 2. Round-Robin Routing +Cycles through workers in order. + +```bash +--policy round_robin +``` + +### 3. Power of Two Choices +Samples two workers and routes to the less loaded one. + +```bash +--policy power_of_two +``` + +### 4. Cache-Aware Load Balancing (Default) + +The most sophisticated policy that combines cache optimization with load balancing: + +```bash +--policy cache_aware \ +--cache-threshold 0.5 \ +--balance-abs-threshold 32 \ +--balance-rel-threshold 1.0001 +``` + +#### How It Works + +1. **Load Assessment**: Checks if the system is balanced + - Imbalanced if: `(max_load - min_load) > balance_abs_threshold` AND `max_load > balance_rel_threshold * min_load` + +2. **Routing Decision**: + - **Balanced System**: Uses cache-aware routing + - Routes to worker with highest prefix match if match > `cache_threshold` + - Otherwise routes to worker with most available cache capacity + - **Imbalanced System**: Uses shortest queue routing to the least busy worker + +3. **Cache Management**: + - Maintains approximate radix trees per worker + - Periodically evicts LRU entries based on `--eviction-interval-secs` and `--max-tree-size` + +### Data Parallelism Aware Routing + +Enables fine-grained control over data parallel replicas: + +```bash +--dp-aware \ +--api-key your_api_key # Required for worker authentication +``` + +This mode coordinates with SGLang's DP controller for optimized request distribution across data parallel ranks. + +## Configuration Reference + +### Core Settings + +| Parameter | Type | Default | Description | +| --------------------------- | ---- | ----------- | --------------------------------------------------------------- | +| `--host` | str | 127.0.0.1 | Router server host address | +| `--port` | int | 30000 | Router server port | +| `--worker-urls` | list | [] | Worker URLs for separate launch mode | +| `--policy` | str | cache_aware | Routing policy (random, round_robin, cache_aware, power_of_two) | +| `--max-concurrent-requests` | int | 64 | Maximum concurrent requests (rate limiting) | +| `--request-timeout-secs` | int | 600 | Request timeout in seconds | +| `--max-payload-size` | int | 256MB | Maximum request payload size | + +### Cache-Aware Routing Parameters + +| Parameter | Type | Default | Description | +| -------------------------- | ----- | -------- | ------------------------------------------------------ | +| `--cache-threshold` | float | 0.5 | Minimum prefix match ratio for cache routing (0.0-1.0) | +| `--balance-abs-threshold` | int | 32 | Absolute load difference threshold | +| `--balance-rel-threshold` | float | 1.0001 | Relative load ratio threshold | +| `--eviction-interval-secs` | int | 60 | Seconds between cache eviction cycles | +| `--max-tree-size` | int | 16777216 | Maximum nodes in routing tree | + +### Fault Tolerance Parameters + +| Parameter | Type | Default | Description | +| ---------------------------- | ----- | ------- | ------------------------------------- | +| `--retry-max-retries` | int | 3 | Maximum retry attempts per request | +| `--retry-initial-backoff-ms` | int | 100 | Initial retry backoff in milliseconds | +| `--retry-max-backoff-ms` | int | 10000 | Maximum retry backoff in milliseconds | +| `--retry-backoff-multiplier` | float | 2.0 | Backoff multiplier between retries | +| `--retry-jitter-factor` | float | 0.1 | Random jitter factor for retries | +| `--disable-retries` | flag | False | Disable retry mechanism | +| `--cb-failure-threshold` | int | 5 | Failures before circuit opens | +| `--cb-success-threshold` | int | 2 | Successes to close circuit | +| `--cb-timeout-duration-secs` | int | 30 | Circuit breaker timeout duration | +| `--cb-window-duration-secs` | int | 60 | Circuit breaker window duration | +| `--disable-circuit-breaker` | flag | False | Disable circuit breaker | + +### Prefill-Decode Disaggregation Parameters + +| Parameter | Type | Default | Description | +| --------------------------------- | ---- | ------- | ----------------------------------------------------- | +| `--pd-disaggregation` | flag | False | Enable PD disaggregated mode | +| `--prefill` | list | [] | Prefill server URLs with optional bootstrap ports | +| `--decode` | list | [] | Decode server URLs | +| `--prefill-policy` | str | None | Routing policy for prefill nodes (overrides --policy) | +| `--decode-policy` | str | None | Routing policy for decode nodes (overrides --policy) | +| `--worker-startup-timeout-secs` | int | 300 | Timeout for worker startup | +| `--worker-startup-check-interval` | int | 10 | Interval between startup checks | + +### Kubernetes Integration + +| Parameter | Type | Default | Description | +| ------------------------------- | ---- | ------------------------ | ---------------------------------------------------- | +| `--service-discovery` | flag | False | Enable Kubernetes service discovery | +| `--selector` | list | [] | Label selector for workers (key1=value1 key2=value2) | +| `--prefill-selector` | list | [] | Label selector for prefill servers in PD mode | +| `--decode-selector` | list | [] | Label selector for decode servers in PD mode | +| `--service-discovery-port` | int | 80 | Port for discovered pods | +| `--service-discovery-namespace` | str | None | Kubernetes namespace to watch | +| `--bootstrap-port-annotation` | str | sglang.ai/bootstrap-port | Annotation for bootstrap ports | + +### Observability + +| Parameter | Type | Default | Description | +| ---------------------- | ---- | --------- | ----------------------------------------------------- | +| `--prometheus-port` | int | 29000 | Prometheus metrics port | +| `--prometheus-host` | str | 127.0.0.1 | Prometheus metrics host | +| `--log-dir` | str | None | Directory for log files | +| `--log-level` | str | info | Logging level (debug, info, warning, error, critical) | +| `--request-id-headers` | list | None | Custom headers for request tracing | + +### CORS Configuration + +| Parameter | Type | Default | Description | +| ------------------------ | ---- | ------- | -------------------- | +| `--cors-allowed-origins` | list | [] | Allowed CORS origins | + +## Advanced Features + +### Kubernetes Service Discovery + +Automatically discover and manage workers in Kubernetes: + +#### Standard Mode +```bash +python -m sglang_router.launch_router \ + --service-discovery \ + --selector app=sglang-worker env=prod \ + --service-discovery-namespace production \ + --service-discovery-port 8000 +``` + +#### Prefill-Decode Disaggregation Mode +```bash +python -m sglang_router.launch_router \ + --pd-disaggregation \ + --service-discovery \ + --prefill-selector app=prefill-server env=prod \ + --decode-selector app=decode-server env=prod \ + --service-discovery-namespace production +``` + +**Note**: The `--bootstrap-port-annotation` (default: `sglang.ai/bootstrap-port`) is used to discover bootstrap ports for prefill servers in PD mode. Prefill pods should have this annotation set to their bootstrap port value. + +### Prometheus Metrics + +Expose metrics for monitoring: + +```bash +python -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 http://worker2:8001 \ + --prometheus-port 29000 \ + --prometheus-host 0.0.0.0 +``` + +Metrics available at `http://localhost:29000/metrics` + +### Request Tracing + +Enable request ID tracking: + +```bash +python -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 http://worker2:8001 \ + --request-id-headers x-request-id x-trace-id +``` + +## Troubleshooting + +### Common Issues + +1. **Workers not connecting**: Ensure workers are fully initialized before starting the router. Use `--worker-startup-timeout-secs` to increase wait time. + +2. **High latency**: Check if cache-aware routing is causing imbalance. Try adjusting `--balance-abs-threshold` and `--balance-rel-threshold`. + +3. **Memory growth**: Reduce `--max-tree-size` or decrease `--eviction-interval-secs` for more aggressive cache cleanup. + +4. **Circuit breaker triggering frequently**: Increase `--cb-failure-threshold` or extend `--cb-window-duration-secs`. + +### Debug Mode + +Enable detailed logging: + +```bash +python -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 http://worker2:8001 \ + --log-level debug \ + --log-dir ./router_logs +``` diff --git a/docs/advanced_features/separate_reasoning.ipynb b/docs/advanced_features/separate_reasoning.ipynb new file mode 100644 index 000000000..0c20c5a08 --- /dev/null +++ b/docs/advanced_features/separate_reasoning.ipynb @@ -0,0 +1,381 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Reasoning Parser\n", + "\n", + "SGLang supports parsing reasoning content out from \"normal\" content for reasoning models such as [DeepSeek R1](https://huggingface.co/deepseek-ai/DeepSeek-R1).\n", + "\n", + "## Supported Models & Parsers\n", + "\n", + "| Model | Reasoning tags | Parser | Notes |\n", + "|---------|-----------------------------|------------------|-------|\n", + "| [DeepSeek‑R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `` … `` | `deepseek-r1` | Supports all variants (R1, R1-0528, R1-Distill) |\n", + "| [DeepSeek‑V3.1](https://huggingface.co/deepseek-ai/DeepSeek-V3.1) | `` … `` | `deepseek-v3` | Supports `thinking` parameter |\n", + "| [Standard Qwen3 models](https://huggingface.co/collections/Qwen/qwen3-67dd247413f0e2e4f653967f) | `` … `` | `qwen3` | Supports `enable_thinking` parameter |\n", + "| [Qwen3-Thinking models](https://huggingface.co/Qwen/Qwen3-235B-A22B-Thinking-2507) | `` … `` | `qwen3` or `qwen3-thinking` | Always generates thinking content |\n", + "| [Kimi models](https://huggingface.co/moonshotai/models) | `◁think▷` … `◁/think▷` | `kimi` | Uses special thinking delimiters |\n", + "| [GPT OSS](https://huggingface.co/openai/gpt-oss-120b) | `<\\|channel\\|>analysis<\\|message\\|>` … `<\\|end\\|>` | `gpt-oss` | N/A |\n", + "### Model-Specific Behaviors\n", + "\n", + "**DeepSeek-R1 Family:**\n", + "- DeepSeek-R1: No `` start tag, jumps directly to thinking content\n", + "- DeepSeek-R1-0528: Generates both `` start and `` end tags\n", + "- Both are handled by the same `deepseek-r1` parser\n", + "\n", + "**DeepSeek-V3 Family:**\n", + "- DeepSeek-V3.1: Hybrid model supporting both thinking and non-thinking modes, use the `deepseek-v3` parser and `thinking` parameter (NOTE: not `enable_thinking`)\n", + "\n", + "**Qwen3 Family:**\n", + "- Standard Qwen3 (e.g., Qwen3-2507): Use `qwen3` parser, supports `enable_thinking` in chat templates\n", + "- Qwen3-Thinking (e.g., Qwen3-235B-A22B-Thinking-2507): Use `qwen3` or `qwen3-thinking` parser, always thinks\n", + "\n", + "**Kimi:**\n", + "- Kimi: Uses special `◁think▷` and `◁/think▷` tags\n", + "\n", + "**GPT OSS:**\n", + "- GPT OSS: Uses special `<|channel|>analysis<|message|>` and `<|end|>` tags" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Usage\n", + "\n", + "### Launching the Server" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Specify the `--reasoning-parser` option." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "from openai import OpenAI\n", + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-7B --host 0.0.0.0 --reasoning-parser deepseek-r1 --log-level warning\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note that `--reasoning-parser` defines the parser used to interpret responses." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### OpenAI Compatible API\n", + "\n", + "Using the OpenAI compatible API, the contract follows the [DeepSeek API design](https://api-docs.deepseek.com/guides/reasoning_model) established with the release of DeepSeek-R1:\n", + "\n", + "- `reasoning_content`: The content of the CoT.\n", + "- `content`: The content of the final answer." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Initialize OpenAI-like client\n", + "client = OpenAI(api_key=\"None\", base_url=f\"http://0.0.0.0:{port}/v1\")\n", + "model_name = client.models.list().data[0].id\n", + "\n", + "messages = [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"What is 1+3?\",\n", + " }\n", + "]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Non-Streaming Request" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response_non_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " stream=False, # Non-streaming\n", + " extra_body={\"separate_reasoning\": True},\n", + ")\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(response_non_stream.choices[0].message.reasoning_content)\n", + "\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(response_non_stream.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Streaming Request" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " stream=True, # Non-streaming\n", + " extra_body={\"separate_reasoning\": True},\n", + ")\n", + "\n", + "reasoning_content = \"\"\n", + "content = \"\"\n", + "for chunk in response_stream:\n", + " if chunk.choices[0].delta.content:\n", + " content += chunk.choices[0].delta.content\n", + " if chunk.choices[0].delta.reasoning_content:\n", + " reasoning_content += chunk.choices[0].delta.reasoning_content\n", + "\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(reasoning_content)\n", + "\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Optionally, you can buffer the reasoning content to the last reasoning chunk (or the first chunk after the reasoning content)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " stream=True, # Non-streaming\n", + " extra_body={\"separate_reasoning\": True, \"stream_reasoning\": False},\n", + ")\n", + "\n", + "reasoning_content = \"\"\n", + "content = \"\"\n", + "for chunk in response_stream:\n", + " if chunk.choices[0].delta.content:\n", + " content += chunk.choices[0].delta.content\n", + " if chunk.choices[0].delta.reasoning_content:\n", + " reasoning_content += chunk.choices[0].delta.reasoning_content\n", + "\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(reasoning_content)\n", + "\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The reasoning separation is enable by default when specify . \n", + "**To disable it, set the `separate_reasoning` option to `False` in request.**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response_non_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " stream=False, # Non-streaming\n", + " extra_body={\"separate_reasoning\": False},\n", + ")\n", + "\n", + "print_highlight(\"==== Original Output ====\")\n", + "print_highlight(response_non_stream.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### SGLang Native API " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import AutoTokenizer\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", + "input = tokenizer.apply_chat_template(\n", + " messages,\n", + " tokenize=False,\n", + " add_generation_prompt=True,\n", + ")\n", + "\n", + "gen_url = f\"http://localhost:{port}/generate\"\n", + "gen_data = {\n", + " \"text\": input,\n", + " \"sampling_params\": {\n", + " \"skip_special_tokens\": False,\n", + " \"max_new_tokens\": 1024,\n", + " \"temperature\": 0.6,\n", + " \"top_p\": 0.95,\n", + " },\n", + "}\n", + "gen_response = requests.post(gen_url, json=gen_data).json()[\"text\"]\n", + "\n", + "print_highlight(\"==== Original Output ====\")\n", + "print_highlight(gen_response)\n", + "\n", + "parse_url = f\"http://localhost:{port}/separate_reasoning\"\n", + "separate_reasoning_data = {\n", + " \"text\": gen_response,\n", + " \"reasoning_parser\": \"deepseek-r1\",\n", + "}\n", + "separate_reasoning_response_json = requests.post(\n", + " parse_url, json=separate_reasoning_data\n", + ").json()\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(separate_reasoning_response_json[\"reasoning_text\"])\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(separate_reasoning_response_json[\"text\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Offline Engine API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sglang as sgl\n", + "from sglang.srt.parser.reasoning_parser import ReasoningParser\n", + "from sglang.utils import print_highlight\n", + "\n", + "llm = sgl.Engine(model_path=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", + "tokenizer = AutoTokenizer.from_pretrained(\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", + "input = tokenizer.apply_chat_template(\n", + " messages,\n", + " tokenize=False,\n", + " add_generation_prompt=True,\n", + ")\n", + "sampling_params = {\n", + " \"max_new_tokens\": 1024,\n", + " \"skip_special_tokens\": False,\n", + " \"temperature\": 0.6,\n", + " \"top_p\": 0.95,\n", + "}\n", + "result = llm.generate(prompt=input, sampling_params=sampling_params)\n", + "\n", + "generated_text = result[\"text\"] # Assume there is only one prompt\n", + "\n", + "print_highlight(\"==== Original Output ====\")\n", + "print_highlight(generated_text)\n", + "\n", + "parser = ReasoningParser(\"deepseek-r1\")\n", + "reasoning_text, text = parser.parse_non_stream(generated_text)\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(reasoning_text)\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(text)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "llm.shutdown()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Supporting New Reasoning Model Schemas\n", + "\n", + "For future reasoning models, you can implement the reasoning parser as a subclass of `BaseReasoningFormatDetector` in `python/sglang/srt/reasoning_parser.py` and specify the reasoning parser for new reasoning model schemas accordingly." + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md new file mode 100644 index 000000000..873fa8b05 --- /dev/null +++ b/docs/advanced_features/server_arguments.md @@ -0,0 +1,318 @@ +# Server Arguments + +This page provides a list of server arguments used in the command line to configure the behavior +and performance of the language model server during deployment. These arguments enable users to +customize key aspects of the server, including model selection, parallelism policies, +memory management, and optimization techniques. +You can find all arguments by `python3 -m sglang.launch_server --help` + +## Common launch commands + +- To enable multi-GPU tensor parallelism, add `--tp 2`. If it reports the error "peer access is not supported between these two devices", add `--enable-p2p-check` to the server launch command. + + ```bash + python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --tp 2 + ``` + +- To enable multi-GPU data parallelism, add `--dp 2`. Data parallelism is better for throughput if there is enough memory. It can also be used together with tensor parallelism. The following command uses 4 GPUs in total. We recommend [SGLang Router](../advanced_features/router.md) for data parallelism. + + ```bash + python -m sglang_router.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --dp 2 --tp 2 + ``` + +- If you see out-of-memory errors during serving, try to reduce the memory usage of the KV cache pool by setting a smaller value of `--mem-fraction-static`. The default value is `0.9`. + + ```bash + python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --mem-fraction-static 0.7 + ``` + +- See [hyperparameter tuning](hyperparameter_tuning.md) on tuning hyperparameters for better performance. +- For docker and Kubernetes runs, you need to set up shared memory which is used for communication between processes. See `--shm-size` for docker and `/dev/shm` size update for Kubernetes manifests. +- If you see out-of-memory errors during prefill for long prompts, try to set a smaller chunked prefill size. + + ```bash + python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --chunked-prefill-size 4096 + ``` + +- To enable `torch.compile` acceleration, add `--enable-torch-compile`. It accelerates small models on small batch sizes. By default, the cache path is located at `/tmp/torchinductor_root`, you can customize it using environment variable `TORCHINDUCTOR_CACHE_DIR`. For more details, please refer to [PyTorch official documentation](https://pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html) and [Enabling cache for torch.compile](https://docs.sglang.ai/backend/hyperparameter_tuning.html#enabling-cache-for-torch-compile). +- To enable torchao quantization, add `--torchao-config int4wo-128`. It supports other [quantization strategies (INT8/FP8)](https://github.com/sgl-project/sglang/blob/v0.3.6/python/sglang/srt/server_args.py#L671) as well. +- To enable fp8 weight quantization, add `--quantization fp8` on a fp16 checkpoint or directly load a fp8 checkpoint without specifying any arguments. +- To enable fp8 kv cache quantization, add `--kv-cache-dtype fp8_e5m2`. +- If the model does not have a chat template in the Hugging Face tokenizer, you can specify a [custom chat template](../references/custom_chat_template.md). +- To run tensor parallelism on multiple nodes, add `--nnodes 2`. If you have two nodes with two GPUs on each node and want to run TP=4, let `sgl-dev-0` be the hostname of the first node and `50000` be an available port, you can use the following commands. If you meet deadlock, please try to add `--disable-cuda-graph` + + ```bash + # Node 0 + python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --tp 4 --dist-init-addr sgl-dev-0:50000 --nnodes 2 --node-rank 0 + + # Node 1 + python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --tp 4 --dist-init-addr sgl-dev-0:50000 --nnodes 2 --node-rank 1 + ``` + +Please consult the documentation below and [server_args.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py) to learn more about the arguments you may provide when launching a server. + +## Model and tokenizer + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--model-path` | The path of the model weights. This can be a local folder or a Hugging Face repo ID. | None | +| `--tokenizer-path` | The path of the tokenizer. | None | +| `--tokenizer-mode` | Tokenizer mode. 'auto' will use the fast tokenizer if available, and 'slow' will always use the slow tokenizer. | auto | +| `--skip-tokenizer-init` | If set, skip init tokenizer and pass input_ids in generate request. | False | +| `--load-format` | The format of the model weights to load. 'auto' will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. 'pt' will load the weights in the pytorch bin format. 'safetensors' will load the weights in the safetensors format. 'npcache' will load the weights in pytorch format and store a numpy cache to speed up the loading. 'dummy' will initialize the weights with random values, which is mainly for profiling. 'gguf' will load the weights in the gguf format. 'bitsandbytes' will load the weights using bitsandbytes quantization. 'layered' loads weights layer by layer so that one can quantize a layer before loading another to make the peak memory envelope smaller. | auto | +| `--trust-remote-code` | Whether or not to allow for custom models defined on the Hub in their own modeling files. | False | +| `--context-length` | The model's maximum context length. Defaults to None (will use the value from the model's config.json instead). | None | +| `--is-embedding` | Whether to use a CausalLM as an embedding model. | False | +| `--enable-multimodal` | Enable the multimodal functionality for the served model. If the model being served is not multimodal, nothing will happen. | None | +| `--revision` | The specific model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. | None | +| `--model-impl` | Which implementation of the model to use. 'auto' will try to use the SGLang implementation if it exists and fall back to the Transformers implementation if no SGLang implementation is available. 'sglang' will use the SGLang model implementation. 'transformers' will use the Transformers model implementation. | auto | + +## HTTP server + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--host` | The host address for the server. | 127.0.0.1 | +| `--port` | The port number for the server. | 30000 | +| `--skip-server-warmup` | If set, skip the server warmup process. | False | +| `--warmups` | Warmup configurations. | None | +| `--nccl-port` | The port for NCCL initialization. | None | + +## Quantization and data type + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--dtype` | Data type for model weights and activations. 'auto' will use FP16 precision for FP32 and FP16 models, and BF16 precision for BF16 models. 'half' for FP16. Recommended for AWQ quantization. 'float16' is the same as 'half'. 'bfloat16' for a balance between precision and range. 'float' is shorthand for FP32 precision. 'float32' for FP32 precision. | auto | +| `--quantization` | The quantization method. | None | +| `--quantization-param-path` | Path to the JSON file containing the KV cache scaling factors. This should generally be supplied, when KV cache dtype is FP8. Otherwise, KV cache scaling factors default to 1.0, which may cause accuracy issues. | None | +| `--kv-cache-dtype` | Data type for kv cache storage. 'auto' will use model data type. 'fp8_e5m2' and 'fp8_e4m3' is supported for CUDA 11.8+. | auto | + +## Memory and scheduling + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--mem-fraction-static` | The fraction of the memory used for static allocation (model weights and KV cache memory pool). Use a smaller value if you see out-of-memory errors. | None | +| `--max-running-requests` | The maximum number of running requests. | None | +| `--max-total-tokens` | The maximum number of tokens in the memory pool. If not specified, it will be automatically calculated based on the memory usage fraction. This option is typically used for development and debugging purposes. | None | +| `--chunked-prefill-size` | The maximum number of tokens in a chunk for the chunked prefill. Setting this to -1 means disabling chunked prefill. | None | +| `--max-prefill-tokens` | The maximum number of tokens in a prefill batch. The real bound will be the maximum of this value and the model's maximum context length. | 16384 | +| `--schedule-policy` | The scheduling policy of the requests. | fcfs | +| `--schedule-conservativeness` | How conservative the schedule policy is. A larger value means more conservative scheduling. Use a larger value if you see requests being retracted frequently. | 1.0 | +| `--cpu-offload-gb` | How many GBs of RAM to reserve for CPU offloading. | 0 | +| `--page-size` | The number of tokens in a page. | 1 | + +## Runtime options + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--device` | The device to use ('cuda', 'xpu', 'hpu', 'npu', 'cpu'). Defaults to auto-detection if not specified. | None | +| `--tp-size` | The tensor parallelism size. | 1 | +| `--pp-size` | The pipeline parallelism size. | 1 | +| `--max-micro-batch-size` | The maximum micro batch size in pipeline parallelism. | None | +| `--stream-interval` | The interval (or buffer size) for streaming in terms of the token length. A smaller value makes streaming smoother, while a larger value makes the throughput higher. | 1 | +| `--stream-output` | Whether to output as a sequence of disjoint segments. | False | +| `--random-seed` | The random seed. | None | +| `--constrained-json-whitespace-pattern` | Regex pattern for syntactic whitespaces allowed in JSON constrained output. For example, to allow the model generate consecutive whitespaces, set the pattern to [\n\t ]*. | None | +| `--watchdog-timeout` | Set watchdog timeout in seconds. If a forward batch takes longer than this, the server will crash to prevent hanging. | 300 | +| `--dist-timeout` | Set timeout for torch.distributed initialization. | None | +| `--download-dir` | Model download directory for huggingface. | None | +| `--base-gpu-id` | The base GPU ID to start allocating GPUs from. Useful when running multiple instances on the same machine. | 0 | +| `--gpu-id-step` | The delta between consecutive GPU IDs that are used. For example, setting it to 2 will use GPU 0,2,4,.... | 1 | +| `--sleep-on-idle` | Reduce CPU usage when sglang is idle. | False | + +## Logging + +| Arguments | Description | Defaults | +|---------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------| +| `--log-level` | The logging level of all loggers. | info | +| `--log-level-http` | The logging level of HTTP server. If not set, reuse --log-level by default. | None | +| `--log-requests` | Log metadata, inputs, outputs of all requests. The verbosity is decided by --log-requests-level. | False | +| `--log-requests-level` | 0: Log metadata (no sampling parameters). 1: Log metadata and sampling parameters. 2: Log metadata, sampling parameters and partial input/output. 3: Log every input/output. | 0 | +| `--show-time-cost` | Show time cost of custom marks. | False | +| `--enable-metrics` | Enable log prometheus metrics. | False | +| `--bucket-time-to-first-token` | The buckets of time to first token, specified as a list of floats. | None | +| `--bucket-inter-token-latency` | The buckets of inter-token latency, specified as a list of floats. | None | +| `--bucket-e2e-request-latency` | The buckets of end-to-end request latency, specified as a list of floats. | None | +| `--collect-tokens-histogram` | Collect prompt/generation tokens histogram. | False | +| `--kv-events-config` | Config in json format for NVIDIA dynamo KV event publishing. Publishing will be enabled if this flag is used. | None | +| `--decode-log-interval` | The log interval of decode batch. | 40 | +| `--enable-request-time-stats-logging` | Enable per request time stats logging. | False | +| `--prompt-tokens-buckets` | The buckets rule of prompt tokens. Supports 3 rule types: 'default' uses predefined buckets; 'tse ' generates two sides exponential distributed buckets (e.g., 'tse 1000 2 8' generates buckets [984.0, 992.0, 996.0, 998.0, 1000.0, 1002.0, 1004.0, 1008.0, 1016.0]).); 'customer ...' uses custom bucket values (e.g., 'customer 10 50 100 500'). | None | +| `--generation-tokens-buckets` | The buckets rule of prompt tokens. Supports 3 rule types: 'default' uses predefined buckets; 'tse ' generates two sides exponential distributed buckets (e.g., 'tse 1000 2 8' generates buckets [984.0, 992.0, 996.0, 998.0, 1000.0, 1002.0, 1004.0, 1008.0, 1016.0]).); 'customer ...' uses custom bucket values (e.g., 'customer 10 50 100 500'). | None | + +## API related + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--api-key` | Set API key of the server. It is also used in the OpenAI API compatible server. | None | +| `--served-model-name` | Override the model name returned by the v1/models endpoint in OpenAI API server. | None | +| `--chat-template` | The buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server. | None | +| `--completion-template` | The buliltin completion template name or the path of the completion template file. This is only used for OpenAI-compatible API server. only for code completion currently. | None | +| `--file-storage-path` | The path of the file storage in backend. | sglang_storage | +| `--enable-cache-report` | Return number of cached tokens in usage.prompt_tokens_details for each openai request. | False | +| `--reasoning-parser` | Specify the parser for reasoning models, supported parsers are: {list(ReasoningParser.DetectorMap.keys())}. | None | +| `--tool-call-parser` | Specify the parser for handling tool-call interactions. Options include: 'qwen25', 'mistral', 'llama3', 'deepseekv3', 'pythonic', 'kimi_k2', 'qwen3_coder', 'glm45', and 'step3'. | None | + +## Data parallelism + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--dp-size` | The data parallelism size. | 1 | +| `--load-balance-method` | The load balancing strategy for data parallelism. Options include: 'round_robin', 'minimum_tokens'. The Minimum Token algorithm can only be used when DP attention is applied. This algorithm performs load balancing based on the real-time token load of the DP workers. | round_robin | + +## Multi-node distributed serving + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--dist-init-addr` | The host address for initializing distributed backend (e.g., `192.168.0.2:25000`). | None | +| `--nnodes` | The number of nodes. | 1 | +| `--node-rank` | The node rank. | 0 | + +## Model override args in JSON + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--json-model-override-args` | A dictionary in JSON string format used to override default model configurations. | {} | +| `--preferred-sampling-params` | json-formatted sampling settings that will be returned in /get_model_info. | None | + +## LoRA + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--enable-lora` | Enable LoRA support for the model. This argument is automatically set to True if `--lora-paths` is provided for backward compatibility. | False | +| `--max-lora-rank` | The maximum LoRA rank that should be supported. If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. This argument is needed when you expect to dynamically load adapters of larger LoRA rank after server startup. | None | +| `--lora-target-modules` | The union set of all target modules where LoRA should be applied (e.g., `q_proj`, `k_proj`, `gate_proj`). If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. This argument is needed when you expect to dynamically load adapters of different target modules after server startup. You can also set it to `all` to enable LoRA for all supported modules. However, enabling LoRA on additional modules introduces a minor performance overhead. If your application is performance-sensitive, we recommend only specifying the modules for which you plan to load adapters. | None | +| `--lora-paths` | The list of LoRA adapters to load. Each adapter must be specified in one of the following formats: | = | JSON with schema {"lora_name":str,"lora_path":str,"pinned":bool} | None | +| `--max-loras-per-batch` | Maximum number of adapters for a running batch, include base-only request. | 8 | +| `--max-loaded-loras` | If specified, it limits the maximum number of LoRA adapters loaded in CPU memory at a time. The value must be greater than or equal to `--max-loras-per-batch`. | None | +| `--lora-backend` | Choose the kernel backend for multi-LoRA serving. | triton | + +## Kernel backend + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--attention-backend` | Choose the kernels for attention layers. | None | +| `--prefill-attention-backend` | (Experimental) This argument specifies the backend for prefill attention computation. Note that this argument has priority over `attention_backend`. | None | +| `--decode-attention-backend` | (Experimental) This argument specifies the backend for decode attention computation. Note that this argument has priority over `attention_backend`. | None | +| `--sampling-backend` | Choose the kernels for sampling layers. | None | +| `--grammar-backend` | Choose the backend for grammar-guided decoding. | None | +| `--mm-attention-backend` | Set multimodal attention backend. | None | + +## Speculative decoding + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--speculative-algorithm` | Speculative algorithm. | None | +| `--speculative-draft-model-path` | The path of the draft model weights. This can be a local folder or a Hugging Face repo ID. | None | +| `--speculative-num-steps` | The number of steps sampled from draft model in Speculative Decoding. | None | +| `--speculative-eagle-topk` | The number of tokens sampled from the draft model in eagle2 each step. | None | +| `--speculative-num-draft-tokens` | The number of tokens sampled from the draft model in Speculative Decoding. | None | +| `--speculative-accept-threshold-single` | Accept a draft token if its probability in the target model is greater than this threshold. | 1.0 | +| `--speculative-accept-threshold-acc` | The accept probability of a draft token is raised from its target probability p to min(1, p / threshold_acc). | 1.0 | +| `--speculative-token-map` | The path of the draft model's small vocab table. | None | +| `--speculative-attention-mode` | Attention backend for speculative decoding operations (both target verify and draft extend). Can be one of 'prefill' (default) or 'decode'. | Prefill | + +## Expert parallelism + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--ep-size` | The expert parallelism size. | 1 | +| `--moe-a2a-backend` | Select the backend for all-to-all communication for expert parallelism. | none | +| `--moe-runner-backend` | Select the runner backend for MoE. | 'triton' | +| `--deepep-mode` | Select the mode when enable DeepEP MoE, could be `normal`, `low_latency` or `auto`. Default is `auto`, which means `low_latency` for decode batch and `normal` for prefill batch. | auto | +| `--ep-num-redundant-experts` | Allocate this number of redundant experts in expert parallel. | 0 | +| `--ep-dispatch-algorithm` | The algorithm to choose ranks for redundant experts in EPLB. | None | +| `--init-expert-location` | Initial location of EP experts. | trivial | +| `--enable-eplb` | Enable EPLB algorithm. | False | +| `--eplb-algorithm` | Chosen EPLB algorithm. | auto | +| `--eplb-rebalance-num-iterations` | Number of iterations to automatically trigger a EPLB re-balance. | 1000 | +| `--eplb-rebalance-layers-per-chunk` | Number of layers to rebalance per forward pass. | None | +| `--expert-distribution-recorder-mode` | Mode of expert distribution recorder. | None | +| `--expert-distribution-recorder-buffer-size` | Circular buffer size of expert distribution recorder. Set to -1 to denote infinite buffer. | None | +| `--enable-expert-distribution-metrics` | Enable logging metrics for expert balancedness. | False | +| `--deepep-config` | Tuned DeepEP config suitable for your own cluster. It can be either a string with JSON content or a file path. | None | +| `--moe-dense-tp-size` | TP size for MoE dense MLP layers. This flag is useful when, with large TP size, there are errors caused by weights in MLP layers having dimension smaller than the min dimension GEMM supports. | None | + +## Hierarchical cache + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--enable-hierarchical-cache` | Enable hierarchical cache. | False | +| `--hicache-ratio` | The ratio of the size of host KV cache memory pool to the size of device pool. | 2.0 | +| `--hicache-size` | The size of the hierarchical cache. | 0 | +| `--hicache-write-policy` | The write policy for hierarchical cache. | write_through | +| `--hicache-io-backend` | The IO backend for hierarchical cache. | | +| `--hicache-storage-backend` | The storage backend for hierarchical cache. | None | + +## Optimization/debug options + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--disable-radix-cache` | Disable RadixAttention for prefix caching. | False | +| `--cuda-graph-max-bs` | Set the maximum batch size for cuda graph. It will extend the cuda graph capture batch size to this value. | None | +| `--cuda-graph-bs` | Set the list of batch sizes for cuda graph. | None | +| `--disable-cuda-graph` | Disable cuda graph. | False | +| `--disable-cuda-graph-padding` | Disable cuda graph when padding is needed. Still uses cuda graph when padding is not needed. | False | +| `--enable-profile-cuda-graph` | Enable profiling of cuda graph capture. | False | +| `--enable-nccl-nvls` | Enable NCCL NVLS for prefill heavy requests when available. | False | +| `--enable-symm-mem` | Enable NCCL symmetric memory for fast collectives. | False | +| `--enable-tokenizer-batch-encode` | Enable batch tokenization for improved performance when processing multiple text inputs. Do not use with image inputs, pre-tokenized input_ids, or input_embeds. | False | +| `--disable-outlines-disk-cache` | Disable disk cache of outlines to avoid possible crashes related to file system or high concurrency. | False | +| `--disable-custom-all-reduce` | Disable the custom all-reduce kernel and fall back to NCCL. | False | +| `--enable-mscclpp` | Enable using mscclpp for small messages for all-reduce kernel and fall back to NCCL. | False | +| `--disable-overlap-schedule` | Disable the overlap scheduler, which overlaps the CPU scheduler with GPU model worker. | False | +| `--enable-mixed-chunk` | Enabling mixing prefill and decode in a batch when using chunked prefill. | False | +| `--enable-dp-attention` | Enabling data parallelism for attention and tensor parallelism for FFN. The dp size should be equal to the tp size. Currently DeepSeek-V2 and Qwen 2/3 MoE models are supported. | False | +| `--enable-dp-lm-head` | Enable vocabulary parallel across the attention TP group to avoid all-gather across DP groups, optimizing performance under DP attention. | False | +| `--enable-two-batch-overlap` | Enabling two micro batches to overlap. | False | +| `--tbo-token-distribution-threshold` | The threshold of token distribution between two batches in micro-batch-overlap, determines whether to two-batch-overlap or two-chunk-overlap. Set to 0 denote disable two-chunk-overlap. | 0.48 | +| `--enable-torch-compile` | Optimize the model with torch.compile. Experimental feature. | False | +| `--torch-compile-max-bs` | Set the maximum batch size when using torch compile. | 32 | +| `--torchao-config` | Optimize the model with torchao. Experimental feature. Current choices are: int8dq, int8wo, int4wo-, fp8wo, fp8dq-per_tensor, fp8dq-per_row. | | +| `--enable-nan-detection` | Enable the NaN detection for debugging purposes. | False | +| `--enable-p2p-check` | Enable P2P check for GPU access, otherwise the p2p access is allowed by default. | False | +| `--triton-attention-reduce-in-fp32` | Cast the intermediate attention results to fp32 to avoid possible crashes related to fp16. This only affects Triton attention kernels. | False | +| `--triton-attention-num-kv-splits` | The number of KV splits in flash decoding Triton kernel. Larger value is better in longer context scenarios. The default value is 8. | 8 | +| `--num-continuous-decode-steps` | Run multiple continuous decoding steps to reduce scheduling overhead. This can potentially increase throughput but may also increase time-to-first-token latency. The default value is 1, meaning only run one decoding step at a time. | 1 | +| `--delete-ckpt-after-loading` | Delete the model checkpoint after loading the model. | False | +| `--enable-memory-saver` | Allow saving memory using release_memory_occupation and resume_memory_occupation. | False | +| `--allow-auto-truncate` | Allow automatically truncating requests that exceed the maximum input length instead of returning an error. | False | +| `--enable-custom-logit-processor` | Enable users to pass custom logit processors to the server (disabled by default for security). | False | +| `--flashinfer-mla-disable-ragged` | Disable ragged processing in Flashinfer MLA. | False | +| `--disable-shared-experts-fusion` | Disable shared experts fusion. | False | +| `--disable-chunked-prefix-cache` | Disable chunked prefix cache. | False | +| `--disable-fast-image-processor` | Disable fast image processor. | False | +| `--enable-return-hidden-states` | Enable returning hidden states. | False | + +## Debug tensor dumps + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--debug-tensor-dump-output-folder` | The output folder for debug tensor dumps. | None | +| `--debug-tensor-dump-input-file` | The input file for debug tensor dumps. | None | +| `--debug-tensor-dump-inject` | Enable injection of debug tensor dumps. | False | +| `--debug-tensor-dump-prefill-only` | Enable prefill-only mode for debug tensor dumps. | False | + +## PD disaggregation + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--disaggregation-mode` | PD disaggregation mode: "null" (not disaggregated), "prefill" (prefill-only), or "decode" (decode-only). | null | +| `--disaggregation-transfer-backend` | The transfer backend for PD disaggregation. | mooncake | +| `--disaggregation-bootstrap-port` | The bootstrap port for PD disaggregation. | 8998 | +| `--disaggregation-decode-tp` | The decode TP for PD disaggregation. | None | +| `--disaggregation-decode-dp` | The decode DP for PD disaggregation. | None | +| `--disaggregation-prefill-pp` | The prefill PP for PD disaggregation. | 1 | + +## Model weight update + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--custom-weight-loader` | Custom weight loader paths. | None | +| `--weight-loader-disable-mmap` | Disable mmap for weight loader. | False | + +## PD-Multiplexing + +| Arguments | Description | Defaults | +|-----------|-------------|----------| +| `--enable-pdmux` | Enable PD-Multiplexing. | False | +| `--sm-group-num` | Number of SM groups for PD-Multiplexing. | 3 | diff --git a/docs/advanced_features/speculative_decoding.ipynb b/docs/advanced_features/speculative_decoding.ipynb new file mode 100644 index 000000000..aa62b897a --- /dev/null +++ b/docs/advanced_features/speculative_decoding.ipynb @@ -0,0 +1,370 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Speculative Decoding\n", + "\n", + "SGLang now provides an EAGLE-based (EAGLE-2/EAGLE-3) speculative decoding option. Our implementation aims to maximize speed and efficiency and is considered to be among the fastest in open-source LLM engines.\n", + "\n", + "### Performance Highlights\n", + "\n", + "Please see below for the huge improvements on throughput for LLaMA-Instruct 3.1 8B tested on MT bench that can be achieved via EAGLE3 decoding.\n", + "For further details please see the [EAGLE3 paper](https://arxiv.org/pdf/2503.01840).\n", + "\n", + "| Method | Throughput (tokens/s) |\n", + "|--------|----------------|\n", + "| SGLang (w/o speculative, 1x H100) | 158.34 tokens/s |\n", + "| SGLang + EAGLE-2 (1x H100) | 244.10 tokens/s |\n", + "| SGLang + EAGLE-3 (1x H100) | 373.25 tokens/s |" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## EAGLE Decoding\n", + "\n", + "To enable EAGLE speculative decoding the following parameters are relevant:\n", + "* `speculative_draft_model_path`: Specifies draft model. This parameter is required.\n", + "* `speculative_num_steps`: Depth of autoregressive drafting. Increases speculation range but risks rejection cascades. Default is 5.\n", + "* `speculative_eagle_topk`: Branching factor per step. Improves candidate diversity, will lead to higher acceptance rate, but more lead to higher memory/compute consumption. Default is 4.\n", + "* `speculative_num_draft_tokens`: Maximum parallel verification capacity. Allows deeper tree evaluation but will lead to higher GPU memory usage. Default is 8.\n", + "\n", + "These parameters are the same for EAGLE-2 and EAGLE-3.\n", + "\n", + "You can find the best combinations of these parameters with [bench_speculative.py](https://github.com/sgl-project/sglang/blob/main/scripts/playground/bench_speculative.py).\n", + "\n", + "In the documentation below, we set `--cuda-graph-max-bs` to be a small value for faster engine startup. For your own workloads, please tune the above parameters together with `--cuda-graph-max-bs`, `--max-running-requests`, `--mem-fraction-static` for the best performance. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE-2 decoding\n", + "\n", + "You can enable EAGLE-2 decoding by setting `--speculative-algorithm EAGLE` and choosing an appropriate model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "import openai" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Llama-2-7b-chat-hf --speculative-algorithm EAGLE \\\n", + " --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B --speculative-num-steps 3 \\\n", + " --speculative-eagle-topk 4 --speculative-num-draft-tokens 16 --cuda-graph-max-bs 8 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Llama-2-7b-chat-hf\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE-2 Decoding with `torch.compile`\n", + "\n", + "You can also enable `torch.compile` for further optimizations and optionally set `--torch-compile-max-bs`:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Llama-2-7b-chat-hf --speculative-algorithm EAGLE \\\n", + " --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B --speculative-num-steps 5 \\\n", + " --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --mem-fraction 0.6 \\\n", + " --enable-torch-compile --torch-compile-max-bs 2 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Llama-2-7b-chat-hf\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE-2 Decoding via Frequency-Ranked Speculative Sampling\n", + "\n", + "By employing a truncated high-frequency token vocabulary in the draft model, Eagle speculative decoding reduces `lm_head` computational overhead while accelerating the pipeline without quality degradation. For more details, checkout [the paper](https://arxiv.org/pdf/arXiv:2502.14856).\n", + "\n", + "In our implementation, set `--speculative-token-map` to enable the optimization. You can get the high-frequency token in FR-Spec from [this model](https://huggingface.co/thunlp/LLaMA3-Instruct-8B-FR-Spec). Or you can obtain high-frequency token by directly downloading these token from [this repo](https://github.com/thunlp/FR-Spec/tree/main?tab=readme-ov-file#prepare-fr-spec-vocabulary-subset).\n", + "\n", + "Thanks for the contribution from [Weilin Zhao](https://github.com/Achazwl) and [Zhousx](https://github.com/Zhou-sx). " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-algorithm EAGLE \\\n", + " --speculative-draft-model-path lmsys/sglang-EAGLE-LLaMA3-Instruct-8B --speculative-num-steps 5 \\\n", + " --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --speculative-token-map thunlp/LLaMA3-Instruct-8B-FR-Spec/freq_32768.pt \\\n", + " --mem-fraction 0.7 --cuda-graph-max-bs 2 --dtype float16 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE-3 Decoding\n", + "\n", + "You can enable EAGLE-3 decoding by setting `--speculative-algorithm EAGLE3` and choosing an appropriate model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B-Instruct --speculative-algorithm EAGLE3 \\\n", + " --speculative-draft-model-path jamesliu1/sglang-EAGLE3-Llama-3.1-Instruct-8B --speculative-num-steps 5 \\\n", + " --speculative-eagle-topk 8 --speculative-num-draft-tokens 32 --mem-fraction 0.6 \\\n", + " --cuda-graph-max-bs 2 --dtype float16 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi Token Prediction\n", + "\n", + "We support [MTP(Multi-Token Prediction)](https://arxiv.org/pdf/2404.19737) in SGLang by using speculative decoding. We use Xiaomi/MiMo-7B-RL model as example here (deepseek mtp usage refer to [deepseek doc](../basic_usage/deepseek.md#multi-token-prediction))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + " python3 -m sglang.launch_server --model-path XiaomiMiMo/MiMo-7B-RL --host 0.0.0.0 --trust-remote-code \\\n", + " --speculative-algorithm EAGLE --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \\\n", + " --mem-fraction 0.5 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "url = f\"http://localhost:{port}/v1/chat/completions\"\n", + "\n", + "data = {\n", + " \"model\": \"XiaomiMiMo/MiMo-7B-RL\",\n", + " \"messages\": [{\"role\": \"user\", \"content\": \"What is the capital of France?\"}],\n", + "}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.json())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## References\n", + "\n", + "EAGLE process is as follows:\n", + "\n", + "- Within EAGLE the draft model predicts the next feature vector, i.e. the last hidden state of the original LLM, using the feature sequence $(f_1, ..., f_k)$ and the token sequence $(t_2, ..., t_{k+1})$. \n", + "- The next token is then sampled from $p_{k+2}=\\text{LMHead}(f_{k+1})$. Afterwards, the two sequences are extended in a tree style—branching out multiple potential continuations, with the branching factor per step controlled by the `speculative_eagle_topk` parameter—to ensure a more coherent connection of context, and are given as input again.\n", + "- EAGLE-2 additionally uses the draft model to evaluate how probable certain branches in the draft tree are, dynamically stopping the expansion of unlikely branches. After the expansion phase, reranking is employed to select only the top `speculative_num_draft_tokens` final nodes as draft tokens.\n", + "- EAGLE-3 removes the feature prediction objective, incorporates low and mid-layer features, and is trained in an on-policy manner.\n", + "\n", + "This enhances drafting accuracy by operating on the features instead of tokens for more regular inputs and passing the tokens from the next timestep additionally to minimize randomness effects from sampling. Furthermore the dynamic adjustment of the draft tree and selection of reranked final nodes increases acceptance rate of draft tokens further. For more details see [EAGLE-2](https://arxiv.org/abs/2406.16858) and [EAGLE-3](https://arxiv.org/abs/2503.01840) paper.\n", + "\n", + "\n", + "For guidance how to train your own EAGLE model please see the [EAGLE repo](https://github.com/SafeAILab/EAGLE/tree/main?tab=readme-ov-file#train)." + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/advanced_features/structured_outputs.ipynb b/docs/advanced_features/structured_outputs.ipynb new file mode 100644 index 000000000..1382f1e0e --- /dev/null +++ b/docs/advanced_features/structured_outputs.ipynb @@ -0,0 +1,853 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Structured Outputs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can specify a JSON schema, [regular expression](https://en.wikipedia.org/wiki/Regular_expression) or [EBNF](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form) to constrain the model output. The model output will be guaranteed to follow the given constraints. Only one constraint parameter (`json_schema`, `regex`, or `ebnf`) can be specified for a request.\n", + "\n", + "SGLang supports three grammar backends:\n", + "\n", + "- [XGrammar](https://github.com/mlc-ai/xgrammar)(default): Supports JSON schema, regular expression, and EBNF constraints.\n", + "- [Outlines](https://github.com/dottxt-ai/outlines): Supports JSON schema and regular expression constraints.\n", + "- [Llguidance](https://github.com/guidance-ai/llguidance): Supports JSON schema, regular expression, and EBNF constraints.\n", + "\n", + "We suggest using XGrammar for its better performance and utility. XGrammar currently uses the [GGML BNF format](https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md). For more details, see [XGrammar technical overview](https://blog.mlc.ai/2024/11/22/achieving-efficient-flexible-portable-structured-generation-with-xgrammar).\n", + "\n", + "To use Outlines, simply add `--grammar-backend outlines` when launching the server.\n", + "To use llguidance, add `--grammar-backend llguidance` when launching the server.\n", + "If no backend is specified, XGrammar will be used as the default.\n", + "\n", + "For better output quality, **It's advisable to explicitly include instructions in the prompt to guide the model to generate the desired format.** For example, you can specify, 'Please generate the output in the following JSON format: ...'.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## OpenAI Compatible API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "import os\n", + "\n", + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", + "\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --host 0.0.0.0 --log-level warning\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### JSON\n", + "\n", + "you can directly define a JSON schema or use [Pydantic](https://docs.pydantic.dev/latest/) to define and validate the response." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Using Pydantic**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pydantic import BaseModel, Field\n", + "\n", + "\n", + "# Define the schema using Pydantic\n", + "class CapitalInfo(BaseModel):\n", + " name: str = Field(..., pattern=r\"^\\w+$\", description=\"Name of the capital city\")\n", + " population: int = Field(..., description=\"Population of the capital city\")\n", + "\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"Please generate the information of the capital of France in the JSON format.\",\n", + " },\n", + " ],\n", + " temperature=0,\n", + " max_tokens=128,\n", + " response_format={\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": {\n", + " \"name\": \"foo\",\n", + " # convert the pydantic model to json schema\n", + " \"schema\": CapitalInfo.model_json_schema(),\n", + " },\n", + " },\n", + ")\n", + "\n", + "response_content = response.choices[0].message.content\n", + "# validate the JSON response by the pydantic model\n", + "capital_info = CapitalInfo.model_validate_json(response_content)\n", + "print_highlight(f\"Validated response: {capital_info.model_dump_json()}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**JSON Schema Directly**\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "json_schema = json.dumps(\n", + " {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"name\": {\"type\": \"string\", \"pattern\": \"^[\\\\w]+$\"},\n", + " \"population\": {\"type\": \"integer\"},\n", + " },\n", + " \"required\": [\"name\", \"population\"],\n", + " }\n", + ")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"Give me the information of the capital of France in the JSON format.\",\n", + " },\n", + " ],\n", + " temperature=0,\n", + " max_tokens=128,\n", + " response_format={\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": {\"name\": \"foo\", \"schema\": json.loads(json_schema)},\n", + " },\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EBNF" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ebnf_grammar = \"\"\"\n", + "root ::= city | description\n", + "city ::= \"London\" | \"Paris\" | \"Berlin\" | \"Rome\"\n", + "description ::= city \" is \" status\n", + "status ::= \"the capital of \" country\n", + "country ::= \"England\" | \"France\" | \"Germany\" | \"Italy\"\n", + "\"\"\"\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": \"You are a helpful geography bot.\"},\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"Give me the information of the capital of France.\",\n", + " },\n", + " ],\n", + " temperature=0,\n", + " max_tokens=32,\n", + " extra_body={\"ebnf\": ebnf_grammar},\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Regular expression" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"What is the capital of France?\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=128,\n", + " extra_body={\"regex\": \"(Paris|London)\"},\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Structural Tag" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tool_get_current_weather = {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_weather\",\n", + " \"description\": \"Get the current weather in a given location\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"city\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The city to find the weather for, e.g. 'San Francisco'\",\n", + " },\n", + " \"state\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"the two-letter abbreviation for the state that the city is\"\n", + " \" in, e.g. 'CA' which would mean 'California'\",\n", + " },\n", + " \"unit\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The unit to fetch the temperature in\",\n", + " \"enum\": [\"celsius\", \"fahrenheit\"],\n", + " },\n", + " },\n", + " \"required\": [\"city\", \"state\", \"unit\"],\n", + " },\n", + " },\n", + "}\n", + "\n", + "tool_get_current_date = {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_date\",\n", + " \"description\": \"Get the current date and time for a given timezone\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"timezone\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The timezone to fetch the current date and time for, e.g. 'America/New_York'\",\n", + " }\n", + " },\n", + " \"required\": [\"timezone\"],\n", + " },\n", + " },\n", + "}\n", + "\n", + "schema_get_current_weather = tool_get_current_weather[\"function\"][\"parameters\"]\n", + "schema_get_current_date = tool_get_current_date[\"function\"][\"parameters\"]\n", + "\n", + "\n", + "def get_messages():\n", + " return [\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": f\"\"\"\n", + "# Tool Instructions\n", + "- Always execute python code in messages that you share.\n", + "- When looking for real time information use relevant functions if available else fallback to brave_search\n", + "You have access to the following functions:\n", + "Use the function 'get_current_weather' to: Get the current weather in a given location\n", + "{tool_get_current_weather[\"function\"]}\n", + "Use the function 'get_current_date' to: Get the current date and time for a given timezone\n", + "{tool_get_current_date[\"function\"]}\n", + "If a you choose to call a function ONLY reply in the following format:\n", + "<{{start_tag}}={{function_name}}>{{parameters}}{{end_tag}}\n", + "where\n", + "start_tag => ` a JSON dict with the function argument name as key and function argument value as value.\n", + "end_tag => ``\n", + "Here is an example,\n", + "{{\"example_name\": \"example_value\"}}\n", + "Reminder:\n", + "- Function calls MUST follow the specified format\n", + "- Required parameters MUST be specified\n", + "- Only call one function at a time\n", + "- Put the entire function call reply on one line\n", + "- Always add your sources when using search results to answer the user query\n", + "You are a helpful assistant.\"\"\",\n", + " },\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"You are in New York. Please get the current date and time, and the weather.\",\n", + " },\n", + " ]\n", + "\n", + "\n", + "messages = get_messages()\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=messages,\n", + " response_format={\n", + " \"type\": \"structural_tag\",\n", + " \"structures\": [\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_weather,\n", + " \"end\": \"\",\n", + " },\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_date,\n", + " \"end\": \"\",\n", + " },\n", + " ],\n", + " \"triggers\": [\"\",\n", + " \"schema\": schema_get_current_weather,\n", + " \"end\": \"\",\n", + " },\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_date,\n", + " \"end\": \"\",\n", + " },\n", + " ],\n", + " \"triggers\": [\"\",\n", + " \"schema\": schema_get_current_weather,\n", + " \"end\": \"\",\n", + " },\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_date,\n", + " \"end\": \"\",\n", + " },\n", + " ],\n", + " \"triggers\": [\"...` to denote reasoning sections, you might want to allow free-form text within these sections while still enforcing grammar constraints on the rest of the output.\n", + "\n", + "SGLang provides a feature to disable grammar restrictions within reasoning sections. This is particularly useful for models that need to perform complex reasoning steps before providing a structured output.\n", + "\n", + "To enable this feature, use the `--reasoning-parser` flag which decide the think_end_token, such as ``, when launching the server. You can also specify the reasoning parser using the `--reasoning-parser` flag.\n", + "\n", + "## Supported Models\n", + "\n", + "Currently, SGLang supports the following reasoning models:\n", + "- [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d): The reasoning content is wrapped with `` and `` tags.\n", + "- [QwQ](https://huggingface.co/Qwen/QwQ-32B): The reasoning content is wrapped with `` and `` tags.\n", + "\n", + "\n", + "## Usage\n", + "\n", + "## OpenAI Compatible API" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Specify the `--grammar-backend`, `--reasoning-parser` option." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "import os\n", + "\n", + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", + "\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-7B --host 0.0.0.0 --reasoning-parser deepseek-r1 --log-level warning\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### JSON\n", + "\n", + "you can directly define a JSON schema or use [Pydantic](https://docs.pydantic.dev/latest/) to define and validate the response." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Using Pydantic**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pydantic import BaseModel, Field\n", + "\n", + "\n", + "# Define the schema using Pydantic\n", + "class CapitalInfo(BaseModel):\n", + " name: str = Field(..., pattern=r\"^\\w+$\", description=\"Name of the capital city\")\n", + " population: int = Field(..., description=\"Population of the capital city\")\n", + "\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", + " messages=[\n", + " {\n", + " \"role\": \"assistant\",\n", + " \"content\": \"Give me the information and population of the capital of France in the JSON format.\",\n", + " },\n", + " ],\n", + " temperature=0,\n", + " max_tokens=2048,\n", + " response_format={\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": {\n", + " \"name\": \"foo\",\n", + " # convert the pydantic model to json schema\n", + " \"schema\": CapitalInfo.model_json_schema(),\n", + " },\n", + " },\n", + ")\n", + "\n", + "print_highlight(\n", + " f\"reasoing_content: {response.choices[0].message.reasoning_content}\\n\\ncontent: {response.choices[0].message.content}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**JSON Schema Directly**\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "json_schema = json.dumps(\n", + " {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"name\": {\"type\": \"string\", \"pattern\": \"^[\\\\w]+$\"},\n", + " \"population\": {\"type\": \"integer\"},\n", + " },\n", + " \"required\": [\"name\", \"population\"],\n", + " }\n", + ")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", + " messages=[\n", + " {\n", + " \"role\": \"assistant\",\n", + " \"content\": \"Give me the information and population of the capital of France in the JSON format.\",\n", + " },\n", + " ],\n", + " temperature=0,\n", + " max_tokens=2048,\n", + " response_format={\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": {\"name\": \"foo\", \"schema\": json.loads(json_schema)},\n", + " },\n", + ")\n", + "\n", + "print_highlight(\n", + " f\"reasoing_content: {response.choices[0].message.reasoning_content}\\n\\ncontent: {response.choices[0].message.content}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EBNF" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ebnf_grammar = \"\"\"\n", + "root ::= city | description\n", + "city ::= \"London\" | \"Paris\" | \"Berlin\" | \"Rome\"\n", + "description ::= city \" is \" status\n", + "status ::= \"the capital of \" country\n", + "country ::= \"England\" | \"France\" | \"Germany\" | \"Italy\"\n", + "\"\"\"\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": \"You are a helpful geography bot.\"},\n", + " {\n", + " \"role\": \"assistant\",\n", + " \"content\": \"Give me the information and population of the capital of France in the JSON format.\",\n", + " },\n", + " ],\n", + " temperature=0,\n", + " max_tokens=2048,\n", + " extra_body={\"ebnf\": ebnf_grammar},\n", + ")\n", + "\n", + "print_highlight(\n", + " f\"reasoing_content: {response.choices[0].message.reasoning_content}\\n\\ncontent: {response.choices[0].message.content}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Regular expression" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = client.chat.completions.create(\n", + " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", + " messages=[\n", + " {\"role\": \"assistant\", \"content\": \"What is the capital of France?\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=2048,\n", + " extra_body={\"regex\": \"(Paris|London)\"},\n", + ")\n", + "\n", + "print_highlight(\n", + " f\"reasoing_content: {response.choices[0].message.reasoning_content}\\n\\ncontent: {response.choices[0].message.content}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Structural Tag" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tool_get_current_weather = {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_weather\",\n", + " \"description\": \"Get the current weather in a given location\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"city\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The city to find the weather for, e.g. 'San Francisco'\",\n", + " },\n", + " \"state\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"the two-letter abbreviation for the state that the city is\"\n", + " \" in, e.g. 'CA' which would mean 'California'\",\n", + " },\n", + " \"unit\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The unit to fetch the temperature in\",\n", + " \"enum\": [\"celsius\", \"fahrenheit\"],\n", + " },\n", + " },\n", + " \"required\": [\"city\", \"state\", \"unit\"],\n", + " },\n", + " },\n", + "}\n", + "\n", + "tool_get_current_date = {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_date\",\n", + " \"description\": \"Get the current date and time for a given timezone\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"timezone\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The timezone to fetch the current date and time for, e.g. 'America/New_York'\",\n", + " }\n", + " },\n", + " \"required\": [\"timezone\"],\n", + " },\n", + " },\n", + "}\n", + "\n", + "schema_get_current_weather = tool_get_current_weather[\"function\"][\"parameters\"]\n", + "schema_get_current_date = tool_get_current_date[\"function\"][\"parameters\"]\n", + "\n", + "\n", + "def get_messages():\n", + " return [\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": f\"\"\"\n", + "# Tool Instructions\n", + "- Always execute python code in messages that you share.\n", + "- When looking for real time information use relevant functions if available else fallback to brave_search\n", + "You have access to the following functions:\n", + "Use the function 'get_current_weather' to: Get the current weather in a given location\n", + "{tool_get_current_weather[\"function\"]}\n", + "Use the function 'get_current_date' to: Get the current date and time for a given timezone\n", + "{tool_get_current_date[\"function\"]}\n", + "If a you choose to call a function ONLY reply in the following format:\n", + "<{{start_tag}}={{function_name}}>{{parameters}}{{end_tag}}\n", + "where\n", + "start_tag => ` a JSON dict with the function argument name as key and function argument value as value.\n", + "end_tag => ``\n", + "Here is an example,\n", + "{{\"example_name\": \"example_value\"}}\n", + "Reminder:\n", + "- Function calls MUST follow the specified format\n", + "- Required parameters MUST be specified\n", + "- Only call one function at a time\n", + "- Put the entire function call reply on one line\n", + "- Always add your sources when using search results to answer the user query\n", + "You are a helpful assistant.\"\"\",\n", + " },\n", + " {\n", + " \"role\": \"assistant\",\n", + " \"content\": \"You are in New York. Please get the current date and time, and the weather.\",\n", + " },\n", + " ]\n", + "\n", + "\n", + "messages = get_messages()\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", + " messages=messages,\n", + " response_format={\n", + " \"type\": \"structural_tag\",\n", + " \"max_new_tokens\": 2048,\n", + " \"structures\": [\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_weather,\n", + " \"end\": \"\",\n", + " },\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_date,\n", + " \"end\": \"\",\n", + " },\n", + " ],\n", + " \"triggers\": [\"\")[0]\n", + "content = response.json()[\"text\"].split(\"\")[1]\n", + "print_highlight(f\"reasoing_content: {reasoing_content}\\n\\ncontent: {content}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**JSON Schema Directly**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "json_schema = json.dumps(\n", + " {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"name\": {\"type\": \"string\", \"pattern\": \"^[\\\\w]+$\"},\n", + " \"population\": {\"type\": \"integer\"},\n", + " },\n", + " \"required\": [\"name\", \"population\"],\n", + " }\n", + ")\n", + "\n", + "# JSON\n", + "text = tokenizer.apply_chat_template(\n", + " messages, tokenize=False, add_generation_prompt=True\n", + ")\n", + "response = requests.post(\n", + " f\"http://localhost:{port}/generate\",\n", + " json={\n", + " \"text\": text,\n", + " \"sampling_params\": {\n", + " \"temperature\": 0,\n", + " \"max_new_tokens\": 2048,\n", + " \"json_schema\": json_schema,\n", + " },\n", + " },\n", + ")\n", + "\n", + "print_highlight(response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EBNF" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = requests.post(\n", + " f\"http://localhost:{port}/generate\",\n", + " json={\n", + " \"text\": \"Give me the information of the capital of France.\",\n", + " \"sampling_params\": {\n", + " \"max_new_tokens\": 2048,\n", + " \"temperature\": 0,\n", + " \"n\": 3,\n", + " \"ebnf\": (\n", + " \"root ::= city | description\\n\"\n", + " 'city ::= \"London\" | \"Paris\" | \"Berlin\" | \"Rome\"\\n'\n", + " 'description ::= city \" is \" status\\n'\n", + " 'status ::= \"the capital of \" country\\n'\n", + " 'country ::= \"England\" | \"France\" | \"Germany\" | \"Italy\"'\n", + " ),\n", + " },\n", + " \"stream\": False,\n", + " \"return_logprob\": False,\n", + " },\n", + ")\n", + "\n", + "print(response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Regular expression" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = requests.post(\n", + " f\"http://localhost:{port}/generate\",\n", + " json={\n", + " \"text\": \"Paris is the capital of\",\n", + " \"sampling_params\": {\n", + " \"temperature\": 0,\n", + " \"max_new_tokens\": 2048,\n", + " \"regex\": \"(France|England)\",\n", + " },\n", + " },\n", + ")\n", + "print(response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Structural Tag" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "text = tokenizer.apply_chat_template(\n", + " messages, tokenize=False, add_generation_prompt=True\n", + ")\n", + "payload = {\n", + " \"text\": text,\n", + " \"sampling_params\": {\n", + " \"max_new_tokens\": 2048,\n", + " \"structural_tag\": json.dumps(\n", + " {\n", + " \"type\": \"structural_tag\",\n", + " \"structures\": [\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_weather,\n", + " \"end\": \"\",\n", + " },\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_date,\n", + " \"end\": \"\",\n", + " },\n", + " ],\n", + " \"triggers\": [\"\",\n", + " \"schema\": schema_get_current_weather,\n", + " \"end\": \"\",\n", + " },\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_date,\n", + " \"end\": \"\",\n", + " },\n", + " ],\n", + " \"triggers\": [\" is not trimmed.\n", + "\n", + "sampling_params = {\n", + " \"max_new_tokens\": 1024,\n", + " \"temperature\": 0,\n", + " \"top_p\": 0.95,\n", + " \"skip_special_tokens\": False,\n", + "}\n", + "\n", + "# 1) Offline generation\n", + "result = llm.generate(input_ids=input_ids, sampling_params=sampling_params)\n", + "generated_text = result[\"text\"] # Assume there is only one prompt\n", + "\n", + "print_highlight(\"=== Offline Engine Output Text ===\")\n", + "print_highlight(generated_text)\n", + "\n", + "\n", + "# 2) Parse using FunctionCallParser\n", + "def convert_dict_to_tool(tool_dict: dict) -> Tool:\n", + " function_dict = tool_dict.get(\"function\", {})\n", + " return Tool(\n", + " type=tool_dict.get(\"type\", \"function\"),\n", + " function=Function(\n", + " name=function_dict.get(\"name\"),\n", + " description=function_dict.get(\"description\"),\n", + " parameters=function_dict.get(\"parameters\"),\n", + " ),\n", + " )\n", + "\n", + "\n", + "tools = [convert_dict_to_tool(raw_tool) for raw_tool in tools]\n", + "\n", + "parser = FunctionCallParser(tools=tools, tool_call_parser=\"qwen25\")\n", + "normal_text, calls = parser.parse_non_stream(generated_text)\n", + "\n", + "print_highlight(\"=== Parsing Result ===\")\n", + "print(\"Normal text portion:\", normal_text)\n", + "print_highlight(\"Function call portion:\")\n", + "for call in calls:\n", + " # call: ToolCallItem\n", + " print_highlight(f\" - tool name: {call.name}\")\n", + " print_highlight(f\" parameters: {call.parameters}\")\n", + "\n", + "# 3) If needed, perform additional logic on the parsed functions, such as automatically calling the corresponding function to obtain a return value, etc." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "llm.shutdown()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Tool Choice Mode\n", + "\n", + "SGLang supports OpenAI's `tool_choice` parameter to control when and which tools the model should call. This feature is implemented using EBNF (Extended Backus-Naur Form) grammar to ensure reliable tool calling behavior.\n", + "\n", + "### Supported Tool Choice Options\n", + "\n", + "- **`tool_choice=\"required\"`**: Forces the model to call at least one tool\n", + "- **`tool_choice={\"type\": \"function\", \"function\": {\"name\": \"specific_function\"}}`**: Forces the model to call a specific function\n", + "\n", + "### Backend Compatibility\n", + "\n", + "Tool choice is fully supported with the **Xgrammar backend**, which is the default grammar backend (`--grammar-backend xgrammar`). However, it may not be fully supported with other backends such as `outlines`.\n", + "\n", + "### Example: Required Tool Choice" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "from sglang.test.doc_patch import launch_server_cmd\n", + "\n", + "# Start a new server session for tool choice examples\n", + "server_process_tool_choice, port_tool_choice = launch_server_cmd(\n", + " \"python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct --tool-call-parser qwen25 --host 0.0.0.0 --log-level warning\"\n", + ")\n", + "wait_for_server(f\"http://localhost:{port_tool_choice}\")\n", + "\n", + "# Initialize client for tool choice examples\n", + "client_tool_choice = OpenAI(\n", + " api_key=\"None\", base_url=f\"http://0.0.0.0:{port_tool_choice}/v1\"\n", + ")\n", + "model_name_tool_choice = client_tool_choice.models.list().data[0].id\n", + "\n", + "# Example with tool_choice=\"required\" - forces the model to call a tool\n", + "messages_required = [\n", + " {\"role\": \"user\", \"content\": \"Hello, what is the capital of France?\"}\n", + "]\n", + "\n", + "# Define tools\n", + "tools = [\n", + " {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_weather\",\n", + " \"description\": \"Get the current weather in a given location\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"city\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The city to find the weather for, e.g. 'San Francisco'\",\n", + " },\n", + " \"unit\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The unit to fetch the temperature in\",\n", + " \"enum\": [\"celsius\", \"fahrenheit\"],\n", + " },\n", + " },\n", + " \"required\": [\"city\", \"unit\"],\n", + " },\n", + " },\n", + " }\n", + "]\n", + "\n", + "response_required = client_tool_choice.chat.completions.create(\n", + " model=model_name_tool_choice,\n", + " messages=messages_required,\n", + " temperature=0,\n", + " max_tokens=1024,\n", + " tools=tools,\n", + " tool_choice=\"required\", # Force the model to call a tool\n", + ")\n", + "\n", + "print_highlight(\"Response with tool_choice='required':\")\n", + "print(\"Content:\", response_required.choices[0].message.content)\n", + "print(\"Tool calls:\", response_required.choices[0].message.tool_calls)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Example: Specific Function Choice\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Example with specific function choice - forces the model to call a specific function\n", + "messages_specific = [\n", + " {\"role\": \"user\", \"content\": \"What are the most attactive places in France?\"}\n", + "]\n", + "\n", + "response_specific = client_tool_choice.chat.completions.create(\n", + " model=model_name_tool_choice,\n", + " messages=messages_specific,\n", + " temperature=0,\n", + " max_tokens=1024,\n", + " tools=tools,\n", + " tool_choice={\n", + " \"type\": \"function\",\n", + " \"function\": {\"name\": \"get_current_weather\"},\n", + " }, # Force the model to call the specific get_current_weather function\n", + ")\n", + "\n", + "print_highlight(\"Response with specific function choice:\")\n", + "print(\"Content:\", response_specific.choices[0].message.content)\n", + "print(\"Tool calls:\", response_specific.choices[0].message.tool_calls)\n", + "\n", + "if response_specific.choices[0].message.tool_calls:\n", + " tool_call = response_specific.choices[0].message.tool_calls[0]\n", + " print_highlight(f\"Called function: {tool_call.function.name}\")\n", + " print_highlight(f\"Arguments: {tool_call.function.arguments}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process_tool_choice)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Pythonic Tool Call Format (Llama-3.2 / Llama-3.3 / Llama-4)\n", + "\n", + "Some Llama models (such as Llama-3.2-1B, Llama-3.2-3B, Llama-3.3-70B, and Llama-4) support a \"pythonic\" tool call format, where the model outputs function calls as Python code, e.g.:\n", + "\n", + "```python\n", + "[get_current_weather(city=\"San Francisco\", state=\"CA\", unit=\"celsius\")]\n", + "```\n", + "\n", + "- The output is a Python list of function calls, with arguments as Python literals (not JSON).\n", + "- Multiple tool calls can be returned in the same list:\n", + "```python\n", + "[get_current_weather(city=\"San Francisco\", state=\"CA\", unit=\"celsius\"),\n", + " get_current_weather(city=\"New York\", state=\"NY\", unit=\"fahrenheit\")]\n", + "```\n", + "\n", + "For more information, refer to Meta’s documentation on [Zero shot function calling](https://github.com/meta-llama/llama-models/blob/main/models/llama4/prompt_format.md#zero-shot-function-calling---system-message).\n", + "\n", + "Note that this feature is still under development on Blackwell.\n", + "\n", + "### How to enable\n", + "- Launch the server with `--tool-call-parser pythonic`\n", + "- You may also specify --chat-template with the improved template for the model (e.g., `--chat-template=examples/chat_template/tool_chat_template_llama4_pythonic.jinja`).\n", + "This is recommended because the model expects a special prompt format to reliably produce valid pythonic tool call outputs. The template ensures that the prompt structure (e.g., special tokens, message boundaries like `<|eom|>`, and function call delimiters) matches what the model was trained or fine-tuned on. If you do not use the correct chat template, tool calling may fail or produce inconsistent results.\n", + "\n", + "#### Forcing Pythonic Tool Call Output Without a Chat Template\n", + "If you don't want to specify a chat template, you must give the model extremely explicit instructions in your messages to enforce pythonic output. For example, for `Llama-3.2-1B-Instruct`, you need:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \" python3 -m sglang.launch_server --model-path meta-llama/Llama-3.2-1B-Instruct --tool-call-parser pythonic --tp 1 --log-level warning\" # llama-3.2-1b-instruct\n", + ")\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "\n", + "tools = [\n", + " {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_weather\",\n", + " \"description\": \"Get the current weather for a given location.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"location\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The name of the city or location.\",\n", + " }\n", + " },\n", + " \"required\": [\"location\"],\n", + " },\n", + " },\n", + " },\n", + " {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_tourist_attractions\",\n", + " \"description\": \"Get a list of top tourist attractions for a given city.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"city\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The name of the city to find attractions for.\",\n", + " }\n", + " },\n", + " \"required\": [\"city\"],\n", + " },\n", + " },\n", + " },\n", + "]\n", + "\n", + "\n", + "def get_messages():\n", + " return [\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": (\n", + " \"You are a travel assistant. \"\n", + " \"When asked to call functions, ALWAYS respond ONLY with a python list of function calls, \"\n", + " \"using this format: [func_name1(param1=value1, param2=value2), func_name2(param=value)]. \"\n", + " \"Do NOT use JSON, do NOT use variables, do NOT use any other format. \"\n", + " \"Here is an example:\\n\"\n", + " '[get_weather(location=\"Paris\"), get_tourist_attractions(city=\"Paris\")]'\n", + " ),\n", + " },\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": (\n", + " \"I'm planning a trip to Tokyo next week. What's the weather like and what are some top tourist attractions? \"\n", + " \"Propose parallel tool calls at once, using the python list of function calls format as shown above.\"\n", + " ),\n", + " },\n", + " ]\n", + "\n", + "\n", + "messages = get_messages()\n", + "\n", + "client = openai.Client(base_url=f\"http://localhost:{port}/v1\", api_key=\"xxxxxx\")\n", + "model_name = client.models.list().data[0].id\n", + "\n", + "\n", + "response_non_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0,\n", + " top_p=0.9,\n", + " stream=False, # Non-streaming\n", + " tools=tools,\n", + ")\n", + "print_highlight(\"Non-stream response:\")\n", + "print_highlight(response_non_stream)\n", + "\n", + "response_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0,\n", + " top_p=0.9,\n", + " stream=True,\n", + " tools=tools,\n", + ")\n", + "texts = \"\"\n", + "tool_calls = []\n", + "name = \"\"\n", + "arguments = \"\"\n", + "\n", + "for chunk in response_stream:\n", + " if chunk.choices[0].delta.content:\n", + " texts += chunk.choices[0].delta.content\n", + " if chunk.choices[0].delta.tool_calls:\n", + " tool_calls.append(chunk.choices[0].delta.tool_calls[0])\n", + "\n", + "print_highlight(\"Streaming Response:\")\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(texts)\n", + "\n", + "print_highlight(\"==== Tool Call ====\")\n", + "for tool_call in tool_calls:\n", + " print_highlight(tool_call)\n", + "\n", + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> **Note:** \n", + "> The model may still default to JSON if it was heavily finetuned on that format. Prompt engineering (including examples) is the only way to increase the chance of pythonic output if you are not using a chat template." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## How to support a new model?\n", + "1. Update the TOOLS_TAG_LIST in sglang/srt/function_call_parser.py with the model’s tool tags. Currently supported tags include:\n", + "```\n", + "\tTOOLS_TAG_LIST = [\n", + "\t “<|plugin|>“,\n", + "\t ““,\n", + "\t “<|python_tag|>“,\n", + "\t “[TOOL_CALLS]”\n", + "\t]\n", + "```\n", + "2. Create a new detector class in sglang/srt/function_call_parser.py that inherits from BaseFormatDetector. The detector should handle the model’s specific function call format. For example:\n", + "```\n", + " class NewModelDetector(BaseFormatDetector):\n", + "```\n", + "3. Add the new detector to the MultiFormatParser class that manages all the format detectors." + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/advanced_features/vlm_query.ipynb b/docs/advanced_features/vlm_query.ipynb new file mode 100644 index 000000000..d9a8ae75d --- /dev/null +++ b/docs/advanced_features/vlm_query.ipynb @@ -0,0 +1,325 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Query Vision Language Model" + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## Querying Qwen-VL" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import nest_asyncio\n", + "\n", + "nest_asyncio.apply() # Run this first.\n", + "\n", + "model_path = \"Qwen/Qwen2.5-VL-3B-Instruct\"\n", + "chat_template = \"qwen2-vl\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "# Lets create a prompt.\n", + "\n", + "from io import BytesIO\n", + "import requests\n", + "from PIL import Image\n", + "\n", + "from sglang.srt.parser.conversation import chat_templates\n", + "\n", + "image = Image.open(\n", + " BytesIO(\n", + " requests.get(\n", + " \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + " ).content\n", + " )\n", + ")\n", + "\n", + "conv = chat_templates[chat_template].copy()\n", + "conv.append_message(conv.roles[0], f\"What's shown here: {conv.image_token}?\")\n", + "conv.append_message(conv.roles[1], \"\")\n", + "conv.image_data = [image]\n", + "\n", + "print(conv.get_prompt())\n", + "image" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "### Query via the offline Engine API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "from sglang import Engine\n", + "\n", + "llm = Engine(\n", + " model_path=model_path, chat_template=chat_template, mem_fraction_static=0.8\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "out = llm.generate(prompt=conv.get_prompt(), image_data=[image])\n", + "print(out[\"text\"])" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "### Query via the offline Engine API, but send precomputed embeddings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "# Compute the image embeddings using Huggingface.\n", + "\n", + "from transformers import AutoProcessor\n", + "from transformers import Qwen2_5_VLForConditionalGeneration\n", + "\n", + "processor = AutoProcessor.from_pretrained(model_path, use_fast=True)\n", + "vision = (\n", + " Qwen2_5_VLForConditionalGeneration.from_pretrained(model_path).eval().visual.cuda()\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "processed_prompt = processor(\n", + " images=[image], text=conv.get_prompt(), return_tensors=\"pt\"\n", + ")\n", + "input_ids = processed_prompt[\"input_ids\"][0].detach().cpu().tolist()\n", + "precomputed_embeddings = vision(\n", + " processed_prompt[\"pixel_values\"].cuda(), processed_prompt[\"image_grid_thw\"].cuda()\n", + ")\n", + "\n", + "mm_item = dict(\n", + " modality=\"IMAGE\",\n", + " image_grid_thw=processed_prompt[\"image_grid_thw\"],\n", + " precomputed_embeddings=precomputed_embeddings,\n", + ")\n", + "out = llm.generate(input_ids=input_ids, image_data=[mm_item])\n", + "print(out[\"text\"])" + ] + }, + { + "cell_type": "markdown", + "id": "10", + "metadata": {}, + "source": [ + "## Querying Llama 4 (Vision)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "import nest_asyncio\n", + "\n", + "nest_asyncio.apply() # Run this first.\n", + "\n", + "model_path = \"meta-llama/Llama-4-Scout-17B-16E-Instruct\"\n", + "chat_template = \"llama-4\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "# Lets create a prompt.\n", + "\n", + "from io import BytesIO\n", + "import requests\n", + "from PIL import Image\n", + "\n", + "from sglang.srt.parser.conversation import chat_templates\n", + "\n", + "image = Image.open(\n", + " BytesIO(\n", + " requests.get(\n", + " \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + " ).content\n", + " )\n", + ")\n", + "\n", + "conv = chat_templates[chat_template].copy()\n", + "conv.append_message(conv.roles[0], f\"What's shown here: {conv.image_token}?\")\n", + "conv.append_message(conv.roles[1], \"\")\n", + "conv.image_data = [image]\n", + "\n", + "print(conv.get_prompt())\n", + "print(f\"Image size: {image.size}\")\n", + "\n", + "image" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "### Query via the offline Engine API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if not is_in_ci():\n", + " from sglang import Engine\n", + "\n", + " llm = Engine(\n", + " model_path=model_path,\n", + " trust_remote_code=True,\n", + " enable_multimodal=True,\n", + " mem_fraction_static=0.8,\n", + " tp_size=4,\n", + " attention_backend=\"fa3\",\n", + " context_length=65536,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "if not is_in_ci():\n", + " out = llm.generate(prompt=conv.get_prompt(), image_data=[image])\n", + " print(out[\"text\"])" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "### Query via the offline Engine API, but send precomputed embeddings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "if not is_in_ci():\n", + " # Compute the image embeddings using Huggingface.\n", + "\n", + " from transformers import AutoProcessor\n", + " from transformers import Llama4ForConditionalGeneration\n", + "\n", + " processor = AutoProcessor.from_pretrained(model_path, use_fast=True)\n", + " model = Llama4ForConditionalGeneration.from_pretrained(\n", + " model_path, torch_dtype=\"auto\"\n", + " ).eval()\n", + " vision = model.vision_model.cuda()\n", + " multi_modal_projector = model.multi_modal_projector.cuda()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "if not is_in_ci():\n", + " processed_prompt = processor(\n", + " images=[image], text=conv.get_prompt(), return_tensors=\"pt\"\n", + " )\n", + " print(f'{processed_prompt[\"pixel_values\"].shape=}')\n", + " input_ids = processed_prompt[\"input_ids\"][0].detach().cpu().tolist()\n", + "\n", + " image_outputs = vision(\n", + " processed_prompt[\"pixel_values\"].to(\"cuda\"), output_hidden_states=False\n", + " )\n", + " image_features = image_outputs.last_hidden_state\n", + " vision_flat = image_features.view(-1, image_features.size(-1))\n", + " precomputed_embeddings = multi_modal_projector(vision_flat)\n", + "\n", + " mm_item = dict(modality=\"IMAGE\", precomputed_embeddings=precomputed_embeddings)\n", + " out = llm.generate(input_ids=input_ids, image_data=[mm_item])\n", + " print(out[\"text\"])" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "custom_cell_magics": "kql", + "encoding": "# -*- coding: utf-8 -*-" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/basic_usage/deepseek.md b/docs/basic_usage/deepseek.md new file mode 100644 index 000000000..8a71696f5 --- /dev/null +++ b/docs/basic_usage/deepseek.md @@ -0,0 +1,228 @@ +# DeepSeek Usage + +SGLang provides many optimizations specifically designed for the DeepSeek models, making it the inference engine recommended by the official [DeepSeek team](https://github.com/deepseek-ai/DeepSeek-V3/tree/main?tab=readme-ov-file#62-inference-with-sglang-recommended) from Day 0. + +This document outlines current optimizations for DeepSeek. +For an overview of the implemented features see the completed [Roadmap](https://github.com/sgl-project/sglang/issues/2591). + +## Launch DeepSeek V3.1/V3/R1 with SGLang + +To run DeepSeek V3.1/V3/R1 models, the recommended settings are as follows: + +| Weight Type | Configuration | +|------------|-------------------| +| **Full precision FP8**
*(recommended)* | 8 x H200 | +| | 8 x MI300X | +| | 2 x 8 x H100/800/20 | +| | Xeon 6980P CPU | +| **Full precision BF16** | 2 x 8 x H200 | +| | 2 x 8 x MI300X | +| | 4 x 8 x H100/800/20 | +| | 4 x 8 x A100/A800 | +| **Quantized weights (AWQ)** | 8 x H100/800/20 | +| | 8 x A100/A800 | +| **Quantized weights (int8)** | 16 x A100/800 | +| | 32 x L40S | +| | Xeon 6980P CPU | +| | 2 x Atlas 800I A3 | + + + +Detailed commands for reference: + +- [8 x H200](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#using-docker-recommended) +- [8 x MI300X](../platforms/amd_gpu.md#running-deepseek-v3) +- [2 x 8 x H200](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h208-nodes) +- [4 x 8 x A100](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-four-a1008-nodes) +- [8 x A100 (AWQ)](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-8-a100a800-with-awq-quantization) +- [16 x A100 (int8)](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-16-a100a800-with-int8-quantization) +- [32 x L40S (int8)](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-32-l40s-with-int8-quantization) +- [Xeon 6980P CPU](../platforms/cpu_server.md#example-running-deepseek-r1) +- [2 x Atlas 800I A3 (int8)](../platforms/ascend_npu.md#running-deepseek-v3) + +### Download Weights +If you encounter errors when starting the server, ensure the weights have finished downloading. It's recommended to download them beforehand or restart multiple times until all weights are downloaded. Please refer to [DeepSeek V3](https://huggingface.co/deepseek-ai/DeepSeek-V3-Base#61-inference-with-deepseek-infer-demo-example-only) official guide to download the weights. + +### Launch with one node of 8 x H200 +Please refer to [the example](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#installation--launch). +**Note that Deepseek V3 is already in FP8**, so we should not run it with any quantization arguments like `--quantization fp8 --kv-cache-dtype fp8_e5m2`. + +### Running examples on Multi-node + +- [Serving with two H20*8 nodes](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h208-nodes). + +- [Serving with two H200*8 nodes and docker](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h2008-nodes-and-docker). + +- [Serving with four A100*8 nodes](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-four-a1008-nodes). + +## Optimizations + +### Multi-head Latent Attention (MLA) Throughput Optimizations + +**Description**: [MLA](https://arxiv.org/pdf/2405.04434) is an innovative attention mechanism introduced by the DeepSeek team, aimed at improving inference efficiency. SGLang has implemented specific optimizations for this, including: + +- **Weight Absorption**: By applying the associative law of matrix multiplication to reorder computation steps, this method balances computation and memory access and improves efficiency in the decoding phase. + +- **MLA Attention Backends**: Currently SGLang supports different optimized MLA attention backends, including [FlashAttention3](https://github.com/Dao-AILab/flash-attention), [Flashinfer](https://docs.flashinfer.ai/api/mla.html), [FlashMLA](https://github.com/deepseek-ai/FlashMLA), [CutlassMLA](https://github.com/sgl-project/sglang/pull/5390), **TRTLLM MLA** (optimized for Blackwell architecture), and [Triton](https://github.com/triton-lang/triton) backends. The default FA3 provides good performance across wide workloads. + +- **FP8 Quantization**: W8A8 FP8 and KV Cache FP8 quantization enables efficient FP8 inference. Additionally, we have implemented Batched Matrix Multiplication (BMM) operator to facilitate FP8 inference in MLA with weight absorption. + +- **CUDA Graph & Torch.compile**: Both MLA and Mixture of Experts (MoE) are compatible with CUDA Graph and Torch.compile, which reduces latency and accelerates decoding speed for small batch sizes. + +- **Chunked Prefix Cache**: Chunked prefix cache optimization can increase throughput by cutting prefix cache into chunks, processing them with multi-head attention and merging their states. Its improvement can be significant when doing chunked prefill on long sequences. Currently this optimization is only available for FlashAttention3 backend. + +Overall, with these optimizations, we have achieved up to **7x** acceleration in output throughput compared to the previous version. + +

+ Multi-head Latent Attention for DeepSeek Series Models +

+ +**Usage**: MLA optimization is enabled by default. For MLA models on Blackwell architecture (e.g., B200), the default backend is FlashInfer. To use the optimized TRTLLM MLA backend for prefill and decode operations, explicitly specify `--attention-backend trtllm_mla`. + +**Reference**: Check [Blog](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#deepseek-multi-head-latent-attention-mla-throughput-optimizations) and [Slides](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/lmsys_1st_meetup_deepseek_mla.pdf) for more details. + +### Data Parallelism Attention + +**Description**: This optimization involves data parallelism (DP) for the MLA attention mechanism of DeepSeek Series Models, which allows for a significant reduction in the KV cache size, enabling larger batch sizes. Each DP worker independently handles different types of batches (prefill, decode, idle), which are then synchronized before and after processing through the Mixture-of-Experts (MoE) layer. If you do not use DP attention, KV cache will be duplicated among all TP ranks. + +

+ Data Parallelism Attention for DeepSeek Series Models +

+ +With data parallelism attention enabled, we have achieved up to **1.9x** decoding throughput improvement compared to the previous version. + +

+ Data Parallelism Attention Performance Comparison +

+ +**Usage**: +- Append `--enable-dp-attention --tp 8 --dp 8` to the server arguments when using 8 H200 GPUs. This optimization improves peak throughput in high batch size scenarios where the server is limited by KV cache capacity. However, it is not recommended for low-latency, small-batch use cases. +- DP and TP attention can be flexibly combined. For example, to deploy DeepSeek-V3/R1 on 2 nodes with 8 H100 GPUs each, you can specify `--enable-dp-attention --tp 16 --dp 2`. This configuration runs attention with 2 DP groups, each containing 8 TP GPUs. + +**Reference**: Check [Blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models). + +### Multi Node Tensor Parallelism + +**Description**: For users with limited memory on a single node, SGLang supports serving DeepSeek Series Models, including DeepSeek V3, across multiple nodes using tensor parallelism. This approach partitions the model parameters across multiple GPUs or nodes to handle models that are too large for one node's memory. + +**Usage**: Check [here](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-2-h208) for usage examples. + +### Block-wise FP8 + +**Description**: SGLang implements block-wise FP8 quantization with two key optimizations: + +- **Activation**: E4M3 format using per-token-per-128-channel sub-vector scales with online casting. + +- **Weight**: Per-128x128-block quantization for better numerical stability. + +- **DeepGEMM**: The [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernel library optimized for FP8 matrix multiplications. + +**Usage**: The activation and weight optimization above are turned on by default for DeepSeek V3 models. DeepGEMM is enabled by default on NVIDIA Hopper GPUs and disabled by default on other devices. DeepGEMM can also be manually turned off by setting the environment variable `SGL_ENABLE_JIT_DEEPGEMM=0`. + +Before serving the DeepSeek model, precompile the DeepGEMM kernels using: +```bash +python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code +``` +The precompilation process typically takes around 10 minutes to complete. + +### Multi-token Prediction +**Description**: SGLang implements DeepSeek V3 Multi-Token Prediction (MTP) based on [EAGLE speculative decoding](https://docs.sglang.ai/advanced_features/speculative_decoding.html#EAGLE-Decoding). With this optimization, the decoding speed can be improved by **1.8x** for batch size 1 and **1.5x** for batch size 32 respectively on H200 TP8 setting. + +**Usage**: +Add arguments `--speculative-algorithm`, `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` to enable this feature. For example: +``` +python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --speculative-algorithm EAGLE --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 --trust-remote-code --tp 8 +``` +- The best configuration for `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` can be searched with [bench_speculative.py](https://github.com/sgl-project/sglang/blob/main/scripts/playground/bench_speculative.py) script for given batch size. The minimum configuration is `--speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2`, which can achieve speedup for larger batch sizes. +- FlashAttention3, FlashMLA, and Triton backend fully supports MTP usage. For FlashInfer backend (`--attention-backend flashinfer`) with speculative decoding,`--speculative-eagle-topk` parameter should be set to `1`. MTP support for the CutlassMLA and TRTLLM MLA backends are still under development. +- To enable DeepSeek MTP for large batch sizes (>32), there are some parameters should be changed (Reference [this discussion](https://github.com/sgl-project/sglang/issues/4543#issuecomment-2737413756)): + - Adjust `--max-running-requests` to a larger number. The default value is `32` for MTP. For larger batch sizes, you should increase this value beyond the default value. + - Set `--cuda-graph-bs`. It's a list of batch sizes for cuda graph capture. The default captured batch sizes for speculative decoding is set [here](https://github.com/sgl-project/sglang/blob/49420741746c8f3e80e0eb17e7d012bfaf25793a/python/sglang/srt/model_executor/cuda_graph_runner.py#L126). You can include more batch sizes into it. + + +### Reasoning Content for DeepSeek R1 & V3.1 + +See [Reasoning Parser](https://docs.sglang.ai/advanced_features/separate_reasoning.html) and [Thinking Parameter for DeepSeek V3.1](https://docs.sglang.ai/basic_usage/openai_api_completions.html#Example:-DeepSeek-V3-Models). + + +### Function calling for DeepSeek Models + +Add arguments `--tool-call-parser deepseekv3` and `--chat-template ./examples/chat_template/tool_chat_template_deepseekv3.jinja`(recommended) to enable this feature. For example (running on 1 * H20 node): + +``` +python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3-0324 --tp 8 --port 30000 --host 0.0.0.0 --mem-fraction-static 0.9 --tool-call-parser deepseekv3 --chat-template ./examples/chat_template/tool_chat_template_deepseekv3.jinja +``` + +Sample Request: + +``` +curl "http://127.0.0.1:30000/v1/chat/completions" \ +-H "Content-Type: application/json" \ +-d '{"temperature": 0, "max_tokens": 100, "model": "deepseek-ai/DeepSeek-V3-0324", "tools": [{"type": "function", "function": {"name": "query_weather", "description": "Get weather of an city, the user should supply a city first", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "The city, e.g. Beijing"}}, "required": ["city"]}}}], "messages": [{"role": "user", "content": "Hows the weather like in Qingdao today"}]}' +``` + +Expected Response + +``` +{"id":"6501ef8e2d874006bf555bc80cddc7c5","object":"chat.completion","created":1745993638,"model":"deepseek-ai/DeepSeek-V3-0324","choices":[{"index":0,"message":{"role":"assistant","content":null,"reasoning_content":null,"tool_calls":[{"id":"0","index":null,"type":"function","function":{"name":"query_weather","arguments":"{\"city\": \"Qingdao\"}"}}]},"logprobs":null,"finish_reason":"tool_calls","matched_stop":null}],"usage":{"prompt_tokens":116,"total_tokens":138,"completion_tokens":22,"prompt_tokens_details":null}} + +``` +Sample Streaming Request: +``` +curl "http://127.0.0.1:30000/v1/chat/completions" \ +-H "Content-Type: application/json" \ +-d '{"temperature": 0, "max_tokens": 100, "model": "deepseek-ai/DeepSeek-V3-0324","stream":true,"tools": [{"type": "function", "function": {"name": "query_weather", "description": "Get weather of an city, the user should supply a city first", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "The city, e.g. Beijing"}}, "required": ["city"]}}}], "messages": [{"role": "user", "content": "Hows the weather like in Qingdao today"}]}' +``` +Expected Streamed Chunks (simplified for clarity): +``` +data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"{\""}}]}}]} +data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"city"}}]}}]} +data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\":\""}}]}}]} +data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"Q"}}]}}]} +data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"ing"}}]}}]} +data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"dao"}}]}}]} +data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\"}"}}]}}]} +data: {"choices":[{"delta":{"tool_calls":null}}], "finish_reason": "tool_calls"} +data: [DONE] +``` +The client needs to concatenate all arguments fragments to reconstruct the complete tool call: +``` +{"city": "Qingdao"} +``` +Important Notes: +1. Use a lower `"temperature"` value for better results. +2. To receive more consistent tool call results, it is recommended to use `--chat-template examples/chat_template/tool_chat_template_deepseekv3.jinja`. It provides an improved unified prompt. + + +## FAQ + +**Q: Model loading is taking too long, and I'm encountering an NCCL timeout. What should I do?** + +A: If you're experiencing extended model loading times and an NCCL timeout, you can try increasing the timeout duration. Add the argument `--dist-timeout 3600` when launching your model. This will set the timeout to one hour, which often resolves the issue. diff --git a/docs/basic_usage/gpt_oss.md b/docs/basic_usage/gpt_oss.md new file mode 100644 index 000000000..240463ec4 --- /dev/null +++ b/docs/basic_usage/gpt_oss.md @@ -0,0 +1,114 @@ +# GPT OSS Usage + +Please refer to [https://github.com/sgl-project/sglang/issues/8833](https://github.com/sgl-project/sglang/issues/8833). + +## Responses API & Built-in Tools + +### Responses API + +GPT‑OSS is compatible with the OpenAI Responses API. Use `client.responses.create(...)` with `model`, `instructions`, `input`, and optional `tools` to enable built‑in tool use. + +### Built-in Tools + +GPT‑OSS can call built‑in tools for web search and Python execution. You can use the demo tool server or connect to external MCP tool servers. + +#### Python Tool + +- Executes short Python snippets for calculations, parsing, and quick scripts. +- By default runs in a Docker-based sandbox. To run on the host, set `PYTHON_EXECUTION_BACKEND=UV` (this executes model-generated code locally; use with care). +- Ensure Docker is available if you are not using the UV backend. It is recommended to run `docker pull python:3.11` in advance. + +#### Web Search Tool + +- Uses the Exa backend for web search. +- Requires an Exa API key; set `EXA_API_KEY` in your environment. Create a key at `https://exa.ai`. + +### Tool & Reasoning Parser + +- We support OpenAI Reasoning and Tool Call parser, as well as our SGLang native api for tool call and reasoning. Refer to [reasoning parser](../advanced_features/separate_reasoning.ipynb) and [tool call parser](../advanced_features/function_calling.ipynb) for more details. + + +## Notes + +- Use **Python 3.12** for the demo tools. And install the required `gpt-oss` packages. +- The default demo integrates the web search tool (Exa backend) and a demo Python interpreter via Docker. +- For search, set `EXA_API_KEY`. For Python execution, either have Docker available or set `PYTHON_EXECUTION_BACKEND=UV`. + +Examples: +```bash +export EXA_API_KEY=YOUR_EXA_KEY +# Optional: run Python tool locally instead of Docker (use with care) +export PYTHON_EXECUTION_BACKEND=UV +``` + +Launch the server with the demo tool server: + +`python3 -m sglang.launch_server --model-path openai/gpt-oss-120b --tool-server demo --tp 2` + +For production usage, sglang can act as an MCP client for multiple services. An [example tool server](https://github.com/openai/gpt-oss/tree/main/gpt-oss-mcp-server) is provided. Start the servers and point sglang to them: +```bash +mcp run -t sse browser_server.py:mcp +mcp run -t sse python_server.py:mcp + +python -m sglang.launch_server ... --tool-server ip-1:port-1,ip-2:port-2 +``` +The URLs should be MCP SSE servers that expose server information and well-documented tools. These tools are added to the system prompt so the model can use them. + +### Quick Demo + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:30000/v1", + api_key="sk-123456" +) + +tools = [ + {"type": "code_interpreter"}, + {"type": "web_search_preview"}, +] + +# Test python tool +response = client.responses.create( + model="openai/gpt-oss-120b", + instructions="You are a helfpul assistant, you could use python tool to execute code.", + input="Use python tool to calculate the sum of 29138749187 and 29138749187", # 58,277,498,374 + tools=tools +) +print("====== test python tool ======") +print(response.output_text) + +# Test browser tool +response = client.responses.create( + model="openai/gpt-oss-120b", + instructions="You are a helfpul assistant, you could use browser to search the web", + input="Search the web for the latest news about Nvidia stock price", + tools=tools +) +print("====== test browser tool ======") +print(response.output_text) +``` + +Example output: +``` +====== test python tool ====== +The sum of 29,138,749,187 and 29,138,749,187 is **58,277,498,374**. +====== test browser tool ====== +**Recent headlines on Nvidia (NVDA) stock** + +| Date (2025) | Source | Key news points | Stock‑price detail | +|-------------|--------|----------------|--------------------| +| **May 13** | Reuters | The market data page shows Nvidia trading “higher” at **$116.61** with no change from the previous close. | **$116.61** – latest trade (delayed ≈ 15 min)【14†L34-L38】 | +| **Aug 18** | CNBC | Morgan Stanley kept an **overweight** rating and lifted its price target to **$206** (up from $200), implying a 14 % upside from the Friday close. The firm notes Nvidia shares have already **jumped 34 % this year**. | No exact price quoted, but the article signals strong upside expectations【9†L27-L31】 | +| **Aug 20** | The Motley Fool | Nvidia is set to release its Q2 earnings on Aug 27. The article lists the **current price of $175.36**, down 0.16 % on the day (as of 3:58 p.m. ET). | **$175.36** – current price on Aug 20【10†L12-L15】【10†L53-L57】 | + +**What the news tells us** + +* Nvidia’s share price has risen sharply this year – up roughly a third according to Morgan Stanley – and analysts are still raising targets (now $206). +* The most recent market quote (Reuters, May 13) was **$116.61**, but the stock has surged since then, reaching **$175.36** by mid‑August. +* Upcoming earnings on **Aug 27** are a focal point; both the Motley Fool and Morgan Stanley expect the results could keep the rally going. + +**Bottom line:** Nvidia’s stock is on a strong upward trajectory in 2025, with price targets climbing toward $200‑$210 and the market price already near $175 as of late August. + +``` diff --git a/docs/basic_usage/llama4.md b/docs/basic_usage/llama4.md new file mode 100644 index 000000000..07cc2b737 --- /dev/null +++ b/docs/basic_usage/llama4.md @@ -0,0 +1,61 @@ +# Llama4 Usage + +[Llama 4](https://github.com/meta-llama/llama-models/blob/main/models/llama4/MODEL_CARD.md) is Meta's latest generation of open-source LLM model with industry-leading performance. + +SGLang has supported Llama 4 Scout (109B) and Llama 4 Maverick (400B) since [v0.4.5](https://github.com/sgl-project/sglang/releases/tag/v0.4.5). + +Ongoing optimizations are tracked in the [Roadmap](https://github.com/sgl-project/sglang/issues/5118). + +## Launch Llama 4 with SGLang + +To serve Llama 4 models on 8xH100/H200 GPUs: + +```bash +python3 -m sglang.launch_server --model-path meta-llama/Llama-4-Scout-17B-16E-Instruct --tp 8 --context-length 1000000 +``` + +### Configuration Tips + +- **OOM Mitigation**: Adjust `--context-length` to avoid a GPU out-of-memory issue. For the Scout model, we recommend setting this value up to 1M on 8\*H100 and up to 2.5M on 8\*H200. For the Maverick model, we don't need to set context length on 8\*H200. When hybrid kv cache is enabled, `--context-length` can be set up to 5M on 8\*H100 and up to 10M on 8\*H200 for the Scout model. + +- **Chat Template**: Add `--chat-template llama-4` for chat completion tasks. +- **Enable Multi-Modal**: Add `--enable-multimodal` for multi-modal capabilities. +- **Enable Hybrid-KVCache**: Add `--hybrid-kvcache-ratio` for hybrid kv cache. Details can be seen in [this PR](https://github.com/sgl-project/sglang/pull/6563) + + +### EAGLE Speculative Decoding +**Description**: SGLang has supported Llama 4 Maverick (400B) with [EAGLE speculative decoding](https://docs.sglang.ai/backend/speculative_decoding.html#EAGLE-Decoding). + +**Usage**: +Add arguments `--speculative-draft-model-path`, `--speculative-algorithm`, `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` to enable this feature. For example: +``` +python3 -m sglang.launch_server --model-path meta-llama/Llama-4-Maverick-17B-128E-Instruct --speculative-algorithm EAGLE3 --speculative-draft-model-path nvidia/Llama-4-Maverick-17B-128E-Eagle3 --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --trust-remote-code --tp 8 --context-length 1000000 +``` + +- **Note** The Llama 4 draft model *nvidia/Llama-4-Maverick-17B-128E-Eagle3* can only recognize conversations in chat mode. + +## Benchmarking Results + +### Accuracy Test with `lm_eval` + +The accuracy on SGLang for both Llama4 Scout and Llama4 Maverick can match the [official benchmark numbers](https://ai.meta.com/blog/llama-4-multimodal-intelligence/). + +Benchmark results on MMLU Pro dataset with 8*H100: +| | Llama-4-Scout-17B-16E-Instruct | Llama-4-Maverick-17B-128E-Instruct | +|--------------------|--------------------------------|-------------------------------------| +| Official Benchmark | 74.3 | 80.5 | +| SGLang | 75.2 | 80.7 | + +Commands: + +```bash +# Llama-4-Scout-17B-16E-Instruct model +python -m sglang.launch_server --model-path meta-llama/Llama-4-Scout-17B-16E-Instruct --port 30000 --tp 8 --mem-fraction-static 0.8 --context-length 65536 +lm_eval --model local-chat-completions --model_args model=meta-llama/Llama-4-Scout-17B-16E-Instruct,base_url=http://localhost:30000/v1/chat/completions,num_concurrent=128,timeout=999999,max_gen_toks=2048 --tasks mmlu_pro --batch_size 128 --apply_chat_template --num_fewshot 0 + +# Llama-4-Maverick-17B-128E-Instruct +python -m sglang.launch_server --model-path meta-llama/Llama-4-Maverick-17B-128E-Instruct --port 30000 --tp 8 --mem-fraction-static 0.8 --context-length 65536 +lm_eval --model local-chat-completions --model_args model=meta-llama/Llama-4-Maverick-17B-128E-Instruct,base_url=http://localhost:30000/v1/chat/completions,num_concurrent=128,timeout=999999,max_gen_toks=2048 --tasks mmlu_pro --batch_size 128 --apply_chat_template --num_fewshot 0 +``` + +Details can be seen in [this PR](https://github.com/sgl-project/sglang/pull/5092). diff --git a/docs/basic_usage/native_api.ipynb b/docs/basic_usage/native_api.ipynb new file mode 100644 index 000000000..5e4ca19a1 --- /dev/null +++ b/docs/basic_usage/native_api.ipynb @@ -0,0 +1,497 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# SGLang Native APIs\n", + "\n", + "Apart from the OpenAI compatible APIs, the SGLang Runtime also provides its native server APIs. We introduce the following APIs:\n", + "\n", + "- `/generate` (text generation model)\n", + "- `/get_model_info`\n", + "- `/get_server_info`\n", + "- `/health`\n", + "- `/health_generate`\n", + "- `/flush_cache`\n", + "- `/update_weights`\n", + "- `/encode`(embedding model)\n", + "- `/v1/rerank`(cross encoder rerank model)\n", + "- `/classify`(reward model)\n", + "- `/start_expert_distribution_record`\n", + "- `/stop_expert_distribution_record`\n", + "- `/dump_expert_distribution_record`\n", + "- A full list of these APIs can be found at [http_server.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/http_server.py)\n", + "\n", + "We mainly use `requests` to test these APIs in the following examples. You can also use `curl`.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct --host 0.0.0.0 --log-level warning\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Generate (text generation model)\n", + "Generate completions. This is similar to the `/v1/completions` in OpenAI API. Detailed parameters can be found in the [sampling parameters](sampling_params.md)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "url = f\"http://localhost:{port}/generate\"\n", + "data = {\"text\": \"What is the capital of France?\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Get Model Info\n", + "\n", + "Get the information of the model.\n", + "\n", + "- `model_path`: The path/name of the model.\n", + "- `is_generation`: Whether the model is used as generation model or embedding model.\n", + "- `tokenizer_path`: The path/name of the tokenizer.\n", + "- `preferred_sampling_params`: The default sampling params specified via `--preferred-sampling-params`. `None` is returned in this example as we did not explicitly configure it in server args.\n", + "- `weight_version`: This field contains the version of the model weights. This is often used to track changes or updates to the model’s trained parameters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://localhost:{port}/get_model_info\"\n", + "\n", + "response = requests.get(url)\n", + "response_json = response.json()\n", + "print_highlight(response_json)\n", + "assert response_json[\"model_path\"] == \"qwen/qwen2.5-0.5b-instruct\"\n", + "assert response_json[\"is_generation\"] is True\n", + "assert response_json[\"tokenizer_path\"] == \"qwen/qwen2.5-0.5b-instruct\"\n", + "assert response_json[\"preferred_sampling_params\"] is None\n", + "assert response_json.keys() == {\n", + " \"model_path\",\n", + " \"is_generation\",\n", + " \"tokenizer_path\",\n", + " \"preferred_sampling_params\",\n", + " \"weight_version\",\n", + "}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Get Server Info\n", + "Gets the server information including CLI arguments, token limits, and memory pool sizes.\n", + "- Note: `get_server_info` merges the following deprecated endpoints:\n", + " - `get_server_args`\n", + " - `get_memory_pool_size` \n", + " - `get_max_total_num_tokens`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://localhost:{port}/get_server_info\"\n", + "\n", + "response = requests.get(url)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Health Check\n", + "- `/health`: Check the health of the server.\n", + "- `/health_generate`: Check the health of the server by generating one token." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://localhost:{port}/health_generate\"\n", + "\n", + "response = requests.get(url)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://localhost:{port}/health\"\n", + "\n", + "response = requests.get(url)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Flush Cache\n", + "\n", + "Flush the radix cache. It will be automatically triggered when the model weights are updated by the `/update_weights` API." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://localhost:{port}/flush_cache\"\n", + "\n", + "response = requests.post(url)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Update Weights From Disk\n", + "\n", + "Update model weights from disk without restarting the server. Only applicable for models with the same architecture and parameter size.\n", + "\n", + "SGLang support `update_weights_from_disk` API for continuous evaluation during training (save checkpoint to disk and update weights from disk).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# successful update with same architecture and size\n", + "\n", + "url = f\"http://localhost:{port}/update_weights_from_disk\"\n", + "data = {\"model_path\": \"qwen/qwen2.5-0.5b-instruct\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.text)\n", + "assert response.json()[\"success\"] is True\n", + "assert response.json()[\"message\"] == \"Succeeded to update model weights.\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# failed update with different parameter size or wrong name\n", + "\n", + "url = f\"http://localhost:{port}/update_weights_from_disk\"\n", + "data = {\"model_path\": \"qwen/qwen2.5-0.5b-instruct-wrong\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "response_json = response.json()\n", + "print_highlight(response_json)\n", + "assert response_json[\"success\"] is False\n", + "assert response_json[\"message\"] == (\n", + " \"Failed to get weights iterator: \"\n", + " \"qwen/qwen2.5-0.5b-instruct-wrong\"\n", + " \" (repository not found).\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Encode (embedding model)\n", + "\n", + "Encode text into embeddings. Note that this API is only available for [embedding models](openai_api_embeddings.ipynb) and will raise an error for generation models.\n", + "Therefore, we launch a new server to server an embedding model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "embedding_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct \\\n", + " --host 0.0.0.0 --is-embedding --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# successful encode for embedding model\n", + "\n", + "url = f\"http://localhost:{port}/encode\"\n", + "data = {\"model\": \"Alibaba-NLP/gte-Qwen2-1.5B-instruct\", \"text\": \"Once upon a time\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "response_json = response.json()\n", + "print_highlight(f\"Text embedding (first 10): {response_json['embedding'][:10]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(embedding_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## v1/rerank (cross encoder rerank model)\n", + "Rerank a list of documents given a query using a cross-encoder model. Note that this API is only available for cross encoder model like [BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3) with `attention-backend` `triton` and `torch_native`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "reranker_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path BAAI/bge-reranker-v2-m3 \\\n", + " --host 0.0.0.0 --disable-radix-cache --chunked-prefill-size -1 --attention-backend triton --is-embedding --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# compute rerank scores for query and documents\n", + "\n", + "url = f\"http://localhost:{port}/v1/rerank\"\n", + "data = {\n", + " \"model\": \"BAAI/bge-reranker-v2-m3\",\n", + " \"query\": \"what is panda?\",\n", + " \"documents\": [\n", + " \"hi\",\n", + " \"The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.\",\n", + " ],\n", + "}\n", + "\n", + "response = requests.post(url, json=data)\n", + "response_json = response.json()\n", + "for item in response_json:\n", + " print_highlight(f\"Score: {item['score']:.2f} - Document: '{item['document']}'\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(reranker_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Classify (reward model)\n", + "\n", + "SGLang Runtime also supports reward models. Here we use a reward model to classify the quality of pairwise generations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Note that SGLang now treats embedding models and reward models as the same type of models.\n", + "# This will be updated in the future.\n", + "\n", + "reward_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path Skywork/Skywork-Reward-Llama-3.1-8B-v0.2 --host 0.0.0.0 --is-embedding --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import AutoTokenizer\n", + "\n", + "PROMPT = (\n", + " \"What is the range of the numeric output of a sigmoid node in a neural network?\"\n", + ")\n", + "\n", + "RESPONSE1 = \"The output of a sigmoid node is bounded between -1 and 1.\"\n", + "RESPONSE2 = \"The output of a sigmoid node is bounded between 0 and 1.\"\n", + "\n", + "CONVS = [\n", + " [{\"role\": \"user\", \"content\": PROMPT}, {\"role\": \"assistant\", \"content\": RESPONSE1}],\n", + " [{\"role\": \"user\", \"content\": PROMPT}, {\"role\": \"assistant\", \"content\": RESPONSE2}],\n", + "]\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(\"Skywork/Skywork-Reward-Llama-3.1-8B-v0.2\")\n", + "prompts = tokenizer.apply_chat_template(CONVS, tokenize=False)\n", + "\n", + "url = f\"http://localhost:{port}/classify\"\n", + "data = {\"model\": \"Skywork/Skywork-Reward-Llama-3.1-8B-v0.2\", \"text\": prompts}\n", + "\n", + "responses = requests.post(url, json=data).json()\n", + "for response in responses:\n", + " print_highlight(f\"reward: {response['embedding'][0]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(reward_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Capture expert selection distribution in MoE models\n", + "\n", + "SGLang Runtime supports recording the number of times an expert is selected in a MoE model run for each expert in the model. This is useful when analyzing the throughput of the model and plan for optimization.\n", + "\n", + "*Note: We only print out the first 10 lines of the csv below for better readability. Please adjust accordingly if you want to analyze the results more deeply.*" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "expert_record_server_process, port = launch_server_cmd(\n", + " \"python3 -m sglang.launch_server --model-path Qwen/Qwen1.5-MoE-A2.7B --host 0.0.0.0 --expert-distribution-recorder-mode stat --log-level warning\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = requests.post(f\"http://localhost:{port}/start_expert_distribution_record\")\n", + "print_highlight(response)\n", + "\n", + "url = f\"http://localhost:{port}/generate\"\n", + "data = {\"text\": \"What is the capital of France?\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.json())\n", + "\n", + "response = requests.post(f\"http://localhost:{port}/stop_expert_distribution_record\")\n", + "print_highlight(response)\n", + "\n", + "response = requests.post(f\"http://localhost:{port}/dump_expert_distribution_record\")\n", + "print_highlight(response)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(expert_record_server_process)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/basic_usage/offline_engine_api.ipynb b/docs/basic_usage/offline_engine_api.ipynb new file mode 100644 index 000000000..9c03e90a7 --- /dev/null +++ b/docs/basic_usage/offline_engine_api.ipynb @@ -0,0 +1,235 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Offline Engine API\n", + "\n", + "SGLang provides a direct inference engine without the need for an HTTP server, especially for use cases where additional HTTP server adds unnecessary complexity or overhead. Here are two general use cases:\n", + "\n", + "- Offline Batch Inference\n", + "- Custom Server on Top of the Engine\n", + "\n", + "This document focuses on the offline batch inference, demonstrating four different inference modes:\n", + "\n", + "- Non-streaming synchronous generation\n", + "- Streaming synchronous generation\n", + "- Non-streaming asynchronous generation\n", + "- Streaming asynchronous generation\n", + "\n", + "Additionally, you can easily build a custom server on top of the SGLang offline engine. A detailed example working in a python script can be found in [custom_server](https://github.com/sgl-project/sglang/blob/main/examples/runtime/engine/custom_server.py).\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Nest Asyncio\n", + "Note that if you want to use **Offline Engine** in ipython or some other nested loop code, you need to add the following code:\n", + "```python\n", + "import nest_asyncio\n", + "\n", + "nest_asyncio.apply()\n", + "\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Advanced Usage\n", + "\n", + "The engine supports [vlm inference](https://github.com/sgl-project/sglang/blob/main/examples/runtime/engine/offline_batch_inference_vlm.py) as well as [extracting hidden states](https://github.com/sgl-project/sglang/blob/main/examples/runtime/hidden_states). \n", + "\n", + "Please see [the examples](https://github.com/sgl-project/sglang/tree/main/examples/runtime/engine) for further use cases." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Offline Batch Inference\n", + "\n", + "SGLang offline engine supports batch inference with efficient scheduling." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# launch the offline engine\n", + "import asyncio\n", + "\n", + "import sglang as sgl\n", + "import sglang.test.doc_patch\n", + "from sglang.utils import async_stream_and_merge, stream_and_merge\n", + "\n", + "llm = sgl.Engine(model_path=\"qwen/qwen2.5-0.5b-instruct\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Non-streaming Synchronous Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompts = [\n", + " \"Hello, my name is\",\n", + " \"The president of the United States is\",\n", + " \"The capital of France is\",\n", + " \"The future of AI is\",\n", + "]\n", + "\n", + "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95}\n", + "\n", + "outputs = llm.generate(prompts, sampling_params)\n", + "for prompt, output in zip(prompts, outputs):\n", + " print(\"===============================\")\n", + " print(f\"Prompt: {prompt}\\nGenerated text: {output['text']}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Streaming Synchronous Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompts = [\n", + " \"Write a short, neutral self-introduction for a fictional character. Hello, my name is\",\n", + " \"Provide a concise factual statement about France’s capital city. The capital of France is\",\n", + " \"Explain possible future trends in artificial intelligence. The future of AI is\",\n", + "]\n", + "\n", + "sampling_params = {\n", + " \"temperature\": 0.2,\n", + " \"top_p\": 0.9,\n", + "}\n", + "\n", + "print(\"\\n=== Testing synchronous streaming generation with overlap removal ===\\n\")\n", + "\n", + "for prompt in prompts:\n", + " print(f\"Prompt: {prompt}\")\n", + " merged_output = stream_and_merge(llm, prompt, sampling_params)\n", + " print(\"Generated text:\", merged_output)\n", + " print()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Non-streaming Asynchronous Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompts = [\n", + " \"Write a short, neutral self-introduction for a fictional character. Hello, my name is\",\n", + " \"Provide a concise factual statement about France’s capital city. The capital of France is\",\n", + " \"Explain possible future trends in artificial intelligence. The future of AI is\",\n", + "]\n", + "\n", + "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95}\n", + "\n", + "print(\"\\n=== Testing asynchronous batch generation ===\")\n", + "\n", + "\n", + "async def main():\n", + " outputs = await llm.async_generate(prompts, sampling_params)\n", + "\n", + " for prompt, output in zip(prompts, outputs):\n", + " print(f\"\\nPrompt: {prompt}\")\n", + " print(f\"Generated text: {output['text']}\")\n", + "\n", + "\n", + "asyncio.run(main())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Streaming Asynchronous Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompts = [\n", + " \"Write a short, neutral self-introduction for a fictional character. Hello, my name is\",\n", + " \"Provide a concise factual statement about France’s capital city. The capital of France is\",\n", + " \"Explain possible future trends in artificial intelligence. The future of AI is\",\n", + "]\n", + "\n", + "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95}\n", + "\n", + "print(\"\\n=== Testing asynchronous streaming generation (no repeats) ===\")\n", + "\n", + "\n", + "async def main():\n", + " for prompt in prompts:\n", + " print(f\"\\nPrompt: {prompt}\")\n", + " print(\"Generated text: \", end=\"\", flush=True)\n", + "\n", + " # Replace direct calls to async_generate with our custom overlap-aware version\n", + " async for cleaned_chunk in async_stream_and_merge(llm, prompt, sampling_params):\n", + " print(cleaned_chunk, end=\"\", flush=True)\n", + "\n", + " print() # New line after each prompt\n", + "\n", + "\n", + "asyncio.run(main())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "llm.shutdown()" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/basic_usage/openai_api.rst b/docs/basic_usage/openai_api.rst new file mode 100644 index 000000000..370abe99c --- /dev/null +++ b/docs/basic_usage/openai_api.rst @@ -0,0 +1,9 @@ +OpenAI-Compatible APIs +====================== + +.. toctree:: + :maxdepth: 1 + + openai_api_completions.ipynb + openai_api_vision.ipynb + openai_api_embeddings.ipynb diff --git a/docs/basic_usage/openai_api_completions.ipynb b/docs/basic_usage/openai_api_completions.ipynb new file mode 100644 index 000000000..6b967709f --- /dev/null +++ b/docs/basic_usage/openai_api_completions.ipynb @@ -0,0 +1,389 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# OpenAI APIs - Completions\n", + "\n", + "SGLang provides OpenAI-compatible APIs to enable a smooth transition from OpenAI services to self-hosted local models.\n", + "A complete reference for the API is available in the [OpenAI API Reference](https://platform.openai.com/docs/api-reference).\n", + "\n", + "This tutorial covers the following popular APIs:\n", + "\n", + "- `chat/completions`\n", + "- `completions`\n", + "\n", + "Check out other tutorials to learn about [vision APIs](openai_api_vision.ipynb) for vision-language models and [embedding APIs](openai_api_embeddings.ipynb) for embedding models." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server\n", + "\n", + "Launch the server in your terminal and wait for it to initialize." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct --host 0.0.0.0 --log-level warning\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "print(f\"Server started on http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Chat Completions\n", + "\n", + "### Usage\n", + "\n", + "The server fully implements the OpenAI API.\n", + "It will automatically apply the chat template specified in the Hugging Face tokenizer, if one is available.\n", + "You can also specify a custom chat template with `--chat-template` when launching the server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"qwen/qwen2.5-0.5b-instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Model Thinking/Reasoning Support\n", + "\n", + "Some models support internal reasoning or thinking processes that can be exposed in the API response. SGLang provides unified support for various reasoning models through the `chat_template_kwargs` parameter and compatible reasoning parsers.\n", + "\n", + "#### Supported Models and Configuration\n", + "\n", + "| Model Family | Chat Template Parameter | Reasoning Parser | Notes |\n", + "|--------------|------------------------|------------------|--------|\n", + "| DeepSeek-R1 (R1, R1-0528, R1-Distill) | `enable_thinking` | `--reasoning-parser deepseek-r1` | Standard reasoning models |\n", + "| DeepSeek-V3.1 | `thinking` | `--reasoning-parser deepseek-v3` | Hybrid model (thinking/non-thinking modes) |\n", + "| Qwen3 (standard) | `enable_thinking` | `--reasoning-parser qwen3` | Hybrid model (thinking/non-thinking modes) |\n", + "| Qwen3-Thinking | N/A (always enabled) | `--reasoning-parser qwen3-thinking` | Always generates reasoning |\n", + "| Kimi | N/A (always enabled) | `--reasoning-parser kimi` | Kimi thinking models |\n", + "| Gpt-Oss | N/A (always enabled) | `--reasoning-parser gpt-oss` | Gpt-Oss thinking models |\n", + "\n", + "#### Basic Usage\n", + "\n", + "To enable reasoning output, you need to:\n", + "1. Launch the server with the appropriate reasoning parser\n", + "2. Set the model-specific parameter in `chat_template_kwargs`\n", + "3. Optionally use `separate_reasoning: False` to not get reasoning content separately (default to `True`)\n", + "\n", + "**Note for Qwen3-Thinking models:** These models always generate thinking content and do not support the `enable_thinking` parameter. Use `--reasoning-parser qwen3-thinking` or `--reasoning-parser qwen3` to parse the thinking content.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Example: Qwen3 Models\n", + "\n", + "```python\n", + "# Launch server:\n", + "# python3 -m sglang.launch_server --model Qwen/Qwen3-4B --reasoning-parser qwen3\n", + "\n", + "from openai import OpenAI\n", + "\n", + "client = OpenAI(\n", + " api_key=\"EMPTY\",\n", + " base_url=f\"http://127.0.0.1:30000/v1\",\n", + ")\n", + "\n", + "model = \"Qwen/Qwen3-4B\"\n", + "messages = [{\"role\": \"user\", \"content\": \"How many r's are in 'strawberry'?\"}]\n", + "\n", + "response = client.chat.completions.create(\n", + " model=model,\n", + " messages=messages,\n", + " extra_body={\n", + " \"chat_template_kwargs\": {\"enable_thinking\": True},\n", + " \"separate_reasoning\": True\n", + " }\n", + ")\n", + "\n", + "print(\"Reasoning:\", response.choices[0].message.reasoning_content)\n", + "print(\"-\"*100)\n", + "print(\"Answer:\", response.choices[0].message.content)\n", + "```\n", + "\n", + "**ExampleOutput:**\n", + "```\n", + "Reasoning: Okay, so the user is asking how many 'r's are in the word 'strawberry'. Let me think. First, I need to make sure I have the word spelled correctly. Strawberry... S-T-R-A-W-B-E-R-R-Y. Wait, is that right? Let me break it down.\n", + "\n", + "Starting with 'strawberry', let's write out the letters one by one. S, T, R, A, W, B, E, R, R, Y. Hmm, wait, that's 10 letters. Let me check again. S (1), T (2), R (3), A (4), W (5), B (6), E (7), R (8), R (9), Y (10). So the letters are S-T-R-A-W-B-E-R-R-Y. \n", + "...\n", + "Therefore, the answer should be three R's in 'strawberry'. But I need to make sure I'm not counting any other letters as R. Let me check again. S, T, R, A, W, B, E, R, R, Y. No other R's. So three in total. Yeah, that seems right.\n", + "\n", + "----------------------------------------------------------------------------------------------------\n", + "Answer: The word \"strawberry\" contains **three** letters 'r'. Here's the breakdown:\n", + "\n", + "1. **S-T-R-A-W-B-E-R-R-Y** \n", + " - The **third letter** is 'R'. \n", + " - The **eighth and ninth letters** are also 'R's. \n", + "\n", + "Thus, the total count is **3**. \n", + "\n", + "**Answer:** 3.\n", + "```\n", + "\n", + "**Note:** Setting `\"enable_thinking\": False` (or omitting it) will result in `reasoning_content` being `None`. Qwen3-Thinking models always generate reasoning content and don't support the `enable_thinking` parameter.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Example: DeepSeek-V3 Models\n", + "\n", + "DeepSeek-V3 models support thinking mode through the `thinking` parameter:\n", + "\n", + "```python\n", + "# Launch server:\n", + "# python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.1 --tp 8 --reasoning-parser deepseek-v3\n", + "\n", + "from openai import OpenAI\n", + "\n", + "client = OpenAI(\n", + " api_key=\"EMPTY\",\n", + " base_url=f\"http://127.0.0.1:30000/v1\",\n", + ")\n", + "\n", + "model = \"deepseek-ai/DeepSeek-V3.1\"\n", + "messages = [{\"role\": \"user\", \"content\": \"How many r's are in 'strawberry'?\"}]\n", + "\n", + "response = client.chat.completions.create(\n", + " model=model,\n", + " messages=messages,\n", + " extra_body={\n", + " \"chat_template_kwargs\": {\"thinking\": True},\n", + " \"separate_reasoning\": True\n", + " }\n", + ")\n", + "\n", + "print(\"Reasoning:\", response.choices[0].message.reasoning_content)\n", + "print(\"-\"*100)\n", + "print(\"Answer:\", response.choices[0].message.content)\n", + "```\n", + "\n", + "**Example Output:**\n", + "```\n", + "Reasoning: First, the question is: \"How many r's are in 'strawberry'?\"\n", + "\n", + "I need to count the number of times the letter 'r' appears in the word \"strawberry\".\n", + "\n", + "Let me write out the word: S-T-R-A-W-B-E-R-R-Y.\n", + "\n", + "Now, I'll go through each letter and count the 'r's.\n", + "...\n", + "So, I have three 'r's in \"strawberry\".\n", + "\n", + "I should double-check. The word is spelled S-T-R-A-W-B-E-R-R-Y. The letters are at positions: 3, 8, and 9 are 'r's. Yes, that's correct.\n", + "\n", + "Therefore, the answer should be 3.\n", + "----------------------------------------------------------------------------------------------------\n", + "Answer: The word \"strawberry\" contains **3** instances of the letter \"r\". Here's a breakdown for clarity:\n", + "\n", + "- The word is spelled: S-T-R-A-W-B-E-R-R-Y\n", + "- The \"r\" appears at the 3rd, 8th, and 9th positions.\n", + "```\n", + "\n", + "**Note:** DeepSeek-V3 models use the `thinking` parameter (not `enable_thinking`) to control reasoning output.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parameters\n", + "\n", + "The chat completions API accepts OpenAI Chat Completions API's parameters. Refer to [OpenAI Chat Completions API](https://platform.openai.com/docs/api-reference/chat/create) for more details.\n", + "\n", + "SGLang extends the standard API with the `extra_body` parameter, allowing for additional customization. One key option within `extra_body` is `chat_template_kwargs`, which can be used to pass arguments to the chat template processor." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = client.chat.completions.create(\n", + " model=\"qwen/qwen2.5-0.5b-instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": \"You are a knowledgeable historian who provides concise responses.\",\n", + " },\n", + " {\"role\": \"user\", \"content\": \"Tell me about ancient Rome\"},\n", + " {\n", + " \"role\": \"assistant\",\n", + " \"content\": \"Ancient Rome was a civilization centered in Italy.\",\n", + " },\n", + " {\"role\": \"user\", \"content\": \"What were their major achievements?\"},\n", + " ],\n", + " temperature=0.3, # Lower temperature for more focused responses\n", + " max_tokens=128, # Reasonable length for a concise response\n", + " top_p=0.95, # Slightly higher for better fluency\n", + " presence_penalty=0.2, # Mild penalty to avoid repetition\n", + " frequency_penalty=0.2, # Mild penalty for more natural language\n", + " n=1, # Single response is usually more stable\n", + " seed=42, # Keep for reproducibility\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Streaming mode is also supported." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "stream = client.chat.completions.create(\n", + " model=\"qwen/qwen2.5-0.5b-instruct\",\n", + " messages=[{\"role\": \"user\", \"content\": \"Say this is a test\"}],\n", + " stream=True,\n", + ")\n", + "for chunk in stream:\n", + " if chunk.choices[0].delta.content is not None:\n", + " print(chunk.choices[0].delta.content, end=\"\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Completions\n", + "\n", + "### Usage\n", + "Completions API is similar to Chat Completions API, but without the `messages` parameter or chat templates." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = client.completions.create(\n", + " model=\"qwen/qwen2.5-0.5b-instruct\",\n", + " prompt=\"List 3 countries and their capitals.\",\n", + " temperature=0,\n", + " max_tokens=64,\n", + " n=1,\n", + " stop=None,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parameters\n", + "\n", + "The completions API accepts OpenAI Completions API's parameters. Refer to [OpenAI Completions API](https://platform.openai.com/docs/api-reference/completions/create) for more details.\n", + "\n", + "Here is an example of a detailed completions request:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = client.completions.create(\n", + " model=\"qwen/qwen2.5-0.5b-instruct\",\n", + " prompt=\"Write a short story about a space explorer.\",\n", + " temperature=0.7, # Moderate temperature for creative writing\n", + " max_tokens=150, # Longer response for a story\n", + " top_p=0.9, # Balanced diversity in word choice\n", + " stop=[\"\\n\\n\", \"THE END\"], # Multiple stop sequences\n", + " presence_penalty=0.3, # Encourage novel elements\n", + " frequency_penalty=0.3, # Reduce repetitive phrases\n", + " n=1, # Generate one completion\n", + " seed=123, # For reproducible results\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Structured Outputs (JSON, Regex, EBNF)\n", + "\n", + "For OpenAI compatible structured outputs API, refer to [Structured Outputs](../advanced_features/structured_outputs.ipynb) for more details.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/basic_usage/openai_api_embeddings.ipynb b/docs/basic_usage/openai_api_embeddings.ipynb new file mode 100644 index 000000000..26e95a4e7 --- /dev/null +++ b/docs/basic_usage/openai_api_embeddings.ipynb @@ -0,0 +1,195 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# OpenAI APIs - Embedding\n", + "\n", + "SGLang provides OpenAI-compatible APIs to enable a smooth transition from OpenAI services to self-hosted local models.\n", + "A complete reference for the API is available in the [OpenAI API Reference](https://platform.openai.com/docs/guides/embeddings).\n", + "\n", + "This tutorial covers the embedding APIs for embedding models. For a list of the supported models see the [corresponding overview page](../supported_models/embedding_models.md)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server\n", + "\n", + "Launch the server in your terminal and wait for it to initialize. Remember to add `--is-embedding` to the command." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "embedding_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct \\\n", + " --host 0.0.0.0 --is-embedding --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using cURL" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess, json\n", + "\n", + "text = \"Once upon a time\"\n", + "\n", + "curl_text = f\"\"\"curl -s http://localhost:{port}/v1/embeddings \\\n", + " -H \"Content-Type: application/json\" \\\n", + " -d '{{\"model\": \"Alibaba-NLP/gte-Qwen2-1.5B-instruct\", \"input\": \"{text}\"}}'\"\"\"\n", + "\n", + "result = subprocess.check_output(curl_text, shell=True)\n", + "\n", + "print(result)\n", + "\n", + "text_embedding = json.loads(result)[\"data\"][0][\"embedding\"]\n", + "\n", + "print_highlight(f\"Text embedding (first 10): {text_embedding[:10]}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Requests" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "text = \"Once upon a time\"\n", + "\n", + "response = requests.post(\n", + " f\"http://localhost:{port}/v1/embeddings\",\n", + " json={\"model\": \"Alibaba-NLP/gte-Qwen2-1.5B-instruct\", \"input\": text},\n", + ")\n", + "\n", + "text_embedding = response.json()[\"data\"][0][\"embedding\"]\n", + "\n", + "print_highlight(f\"Text embedding (first 10): {text_embedding[:10]}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using OpenAI Python Client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "# Text embedding example\n", + "response = client.embeddings.create(\n", + " model=\"Alibaba-NLP/gte-Qwen2-1.5B-instruct\",\n", + " input=text,\n", + ")\n", + "\n", + "embedding = response.data[0].embedding[:10]\n", + "print_highlight(f\"Text embedding (first 10): {embedding}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Input IDs\n", + "\n", + "SGLang also supports `input_ids` as input to get the embedding." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\n", + "from transformers import AutoTokenizer\n", + "\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(\"Alibaba-NLP/gte-Qwen2-1.5B-instruct\")\n", + "input_ids = tokenizer.encode(text)\n", + "\n", + "curl_ids = f\"\"\"curl -s http://localhost:{port}/v1/embeddings \\\n", + " -H \"Content-Type: application/json\" \\\n", + " -d '{{\"model\": \"Alibaba-NLP/gte-Qwen2-1.5B-instruct\", \"input\": {json.dumps(input_ids)}}}'\"\"\"\n", + "\n", + "input_ids_embedding = json.loads(subprocess.check_output(curl_ids, shell=True))[\"data\"][\n", + " 0\n", + "][\"embedding\"]\n", + "\n", + "print_highlight(f\"Input IDs embedding (first 10): {input_ids_embedding[:10]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(embedding_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi-Modal Embedding Model\n", + "Please refer to [Multi-Modal Embedding Model](../supported_models/embedding_models.md)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/basic_usage/openai_api_vision.ipynb b/docs/basic_usage/openai_api_vision.ipynb new file mode 100644 index 000000000..88d1ef7dd --- /dev/null +++ b/docs/basic_usage/openai_api_vision.ipynb @@ -0,0 +1,254 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# OpenAI APIs - Vision\n", + "\n", + "SGLang provides OpenAI-compatible APIs to enable a smooth transition from OpenAI services to self-hosted local models.\n", + "A complete reference for the API is available in the [OpenAI API Reference](https://platform.openai.com/docs/guides/vision).\n", + "This tutorial covers the vision APIs for vision language models.\n", + "\n", + "SGLang supports various vision language models such as Llama 3.2, LLaVA-OneVision, Qwen2.5-VL, Gemma3 and [more](../supported_models/multimodal_language_models.md).\n", + "\n", + "As an alternative to the OpenAI API, you can also use the [SGLang offline engine](https://github.com/sgl-project/sglang/blob/main/examples/runtime/engine/offline_batch_inference_vlm.py)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server\n", + "\n", + "Launch the server in your terminal and wait for it to initialize." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "vision_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-VL-7B-Instruct --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using cURL\n", + "\n", + "Once the server is up, you can send test requests using curl or requests." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess\n", + "\n", + "curl_command = f\"\"\"\n", + "curl -s http://localhost:{port}/v1/chat/completions \\\\\n", + " -H \"Content-Type: application/json\" \\\\\n", + " -d '{{\n", + " \"model\": \"Qwen/Qwen2.5-VL-7B-Instruct\",\n", + " \"messages\": [\n", + " {{\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {{\n", + " \"type\": \"text\",\n", + " \"text\": \"What’s in this image?\"\n", + " }},\n", + " {{\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {{\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + " }}\n", + " }}\n", + " ]\n", + " }}\n", + " ],\n", + " \"max_tokens\": 300\n", + " }}'\n", + "\"\"\"\n", + "\n", + "response = subprocess.check_output(curl_command, shell=True).decode()\n", + "print_highlight(response)\n", + "\n", + "\n", + "response = subprocess.check_output(curl_command, shell=True).decode()\n", + "print_highlight(response)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Requests" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "url = f\"http://localhost:{port}/v1/chat/completions\"\n", + "\n", + "data = {\n", + " \"model\": \"Qwen/Qwen2.5-VL-7B-Instruct\",\n", + " \"messages\": [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"text\", \"text\": \"What’s in this image?\"},\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + " },\n", + " },\n", + " ],\n", + " }\n", + " ],\n", + " \"max_tokens\": 300,\n", + "}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using OpenAI Python Client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "\n", + "client = OpenAI(base_url=f\"http://localhost:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"Qwen/Qwen2.5-VL-7B-Instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\n", + " \"type\": \"text\",\n", + " \"text\": \"What is in this image?\",\n", + " },\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + " },\n", + " },\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=300,\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multiple-Image Inputs\n", + "\n", + "The server also supports multiple images and interleaved text and images if the model supports it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "\n", + "client = OpenAI(base_url=f\"http://localhost:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"Qwen/Qwen2.5-VL-7B-Instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\",\n", + " },\n", + " },\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png\",\n", + " },\n", + " },\n", + " {\n", + " \"type\": \"text\",\n", + " \"text\": \"I have two very different images. They are not related at all. \"\n", + " \"Please describe the first image in one sentence, and then describe the second image in another sentence.\",\n", + " },\n", + " ],\n", + " }\n", + " ],\n", + " temperature=0,\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(vision_process)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/basic_usage/qwen3.md b/docs/basic_usage/qwen3.md new file mode 100644 index 000000000..89295be60 --- /dev/null +++ b/docs/basic_usage/qwen3.md @@ -0,0 +1,27 @@ +# Qwen3-Next Usage + +SGLang has supported Qwen3-Next-80B-A3B-Instruct and Qwen3-Next-80B-A3B-Thinking since [this PR](https://github.com/sgl-project/sglang/pull/10233). + +## Launch Qwen3-Next with SGLang + +To serve Qwen3-Next models on 4xH100/H200 GPUs: + +```bash +python3 -m sglang.launch_server --model Qwen/Qwen3-Next-80B-A3B-Instruct --tp 4 +``` + +### Configuration Tips +- `--max-mamba-cache-size`: Adjust `--max-mamba-cache-size` to increase mamba cache space and max running requests capability. It will decrease KV cache space as a trade-off. You can adjust it according to workload. +- `--mamba-ssm-dtype`: `bfloat16` or `float32`, use `bfloat16` to save mamba cache size and `float32` to get more accurate results. The default setting is `float32`. + +### EAGLE Speculative Decoding +**Description**: SGLang has supported Qwen3-Next models with [EAGLE speculative decoding](https://docs.sglang.ai/advanced_features/speculative_decoding.html#EAGLE-Decoding). + +**Usage**: +Add arguments `--speculative-algorithm`, `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` to enable this feature. For example: + +``` bash +python3 -m sglang.launch_server --model Qwen/Qwen3-Next-80B-A3B-Instruct --tp 4 --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --speculative-algo NEXTN +``` + +Details can be seen in [this PR](https://github.com/sgl-project/sglang/pull/10233). diff --git a/docs/basic_usage/sampling_params.md b/docs/basic_usage/sampling_params.md new file mode 100644 index 000000000..c1394a9fd --- /dev/null +++ b/docs/basic_usage/sampling_params.md @@ -0,0 +1,305 @@ +# Sampling Parameters + +This doc describes the sampling parameters of the SGLang Runtime. It is the low-level endpoint of the runtime. +If you want a high-level endpoint that can automatically handle chat templates, consider using the [OpenAI Compatible API](openai_api_completions.ipynb). + +## `/generate` Endpoint + +The `/generate` endpoint accepts the following parameters in JSON format. For detailed usage, see the [native API doc](native_api.ipynb). The object is defined at `io_struct.py::GenerateReqInput`. You can also read the source code to find more arguments and docs. + +| Argument | Type/Default | Description | +|----------------------------|------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------| +| text | `Optional[Union[List[str], str]] = None` | The input prompt. Can be a single prompt or a batch of prompts. | +| input_ids | `Optional[Union[List[List[int]], List[int]]] = None` | The token IDs for text; one can specify either text or input_ids. | +| input_embeds | `Optional[Union[List[List[List[float]]], List[List[float]]]] = None` | The embeddings for input_ids; one can specify either text, input_ids, or input_embeds. | +| image_data | `Optional[Union[List[List[ImageDataItem]], List[ImageDataItem], ImageDataItem]] = None` | The image input. Can be an image instance, file name, URL, or base64 encoded string. Can be a single image, list of images, or list of lists of images. | +| audio_data | `Optional[Union[List[AudioDataItem], AudioDataItem]] = None` | The audio input. Can be a file name, URL, or base64 encoded string. | +| sampling_params | `Optional[Union[List[Dict], Dict]] = None` | The sampling parameters as described in the sections below. | +| rid | `Optional[Union[List[str], str]] = None` | The request ID. | +| return_logprob | `Optional[Union[List[bool], bool]] = None` | Whether to return log probabilities for tokens. | +| logprob_start_len | `Optional[Union[List[int], int]] = None` | If return_logprob, the start location in the prompt for returning logprobs. Default is "-1", which returns logprobs for output tokens only. | +| top_logprobs_num | `Optional[Union[List[int], int]] = None` | If return_logprob, the number of top logprobs to return at each position. | +| token_ids_logprob | `Optional[Union[List[List[int]], List[int]]] = None` | If return_logprob, the token IDs to return logprob for. | +| return_text_in_logprobs | `bool = False` | Whether to detokenize tokens in text in the returned logprobs. | +| stream | `bool = False` | Whether to stream output. | +| lora_path | `Optional[Union[List[Optional[str]], Optional[str]]] = None` | The path to the LoRA. | +| custom_logit_processor | `Optional[Union[List[Optional[str]], str]] = None` | Custom logit processor for advanced sampling control. Must be a serialized instance of `CustomLogitProcessor` using its `to_str()` method. For usage see below. | +| return_hidden_states | `Union[List[bool], bool] = False` | Whether to return hidden states. | + +## Sampling parameters + +The object is defined at `sampling_params.py::SamplingParams`. You can also read the source code to find more arguments and docs. + +### Core parameters + +| Argument | Type/Default | Description | +|-----------------|----------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| +| max_new_tokens | `int = 128` | The maximum output length measured in tokens. | +| stop | `Optional[Union[str, List[str]]] = None` | One or multiple [stop words](https://platform.openai.com/docs/api-reference/chat/create#chat-create-stop). Generation will stop if one of these words is sampled. | +| stop_token_ids | `Optional[List[int]] = None` | Provide stop words in the form of token IDs. Generation will stop if one of these token IDs is sampled. | +| temperature | `float = 1.0` | [Temperature](https://platform.openai.com/docs/api-reference/chat/create#chat-create-temperature) when sampling the next token. `temperature = 0` corresponds to greedy sampling, a higher temperature leads to more diversity. | +| top_p | `float = 1.0` | [Top-p](https://platform.openai.com/docs/api-reference/chat/create#chat-create-top_p) selects tokens from the smallest sorted set whose cumulative probability exceeds `top_p`. When `top_p = 1`, this reduces to unrestricted sampling from all tokens. | +| top_k | `int = -1` | [Top-k](https://developer.nvidia.com/blog/how-to-get-better-outputs-from-your-large-language-model/#predictability_vs_creativity) randomly selects from the `k` highest-probability tokens. | +| min_p | `float = 0.0` | [Min-p](https://github.com/huggingface/transformers/issues/27670) samples from tokens with probability larger than `min_p * highest_token_probability`. | + +### Penalizers + +| Argument | Type/Default | Description | +|--------------------|------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| +| frequency_penalty | `float = 0.0` | Penalizes tokens based on their frequency in generation so far. Must be between `-2` and `2` where negative numbers encourage repeatment of tokens and positive number encourages sampling of new tokens. The scaling of penalization grows linearly with each appearance of a token. | +| presence_penalty | `float = 0.0` | Penalizes tokens if they appeared in the generation so far. Must be between `-2` and `2` where negative numbers encourage repeatment of tokens and positive number encourages sampling of new tokens. The scaling of the penalization is constant if a token occurred. | +| min_new_tokens | `int = 0` | Forces the model to generate at least `min_new_tokens` until a stop word or EOS token is sampled. Note that this might lead to unintended behavior, for example, if the distribution is highly skewed towards these tokens. | + +### Constrained decoding + +Please refer to our dedicated guide on [constrained decoding](../advanced_features/structured_outputs.ipynb) for the following parameters. + +| Argument | Type/Default | Description | +|-----------------|---------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| +| json_schema | `Optional[str] = None` | JSON schema for structured outputs. | +| regex | `Optional[str] = None` | Regex for structured outputs. | +| ebnf | `Optional[str] = None` | EBNF for structured outputs. | +| structural_tag | `Optional[str] = None` | The structal tag for structured outputs. | + +### Other options + +| Argument | Type/Default | Description | +|-------------------------------|---------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| +| n | `int = 1` | Specifies the number of output sequences to generate per request. (Generating multiple outputs in one request (n > 1) is discouraged; repeating the same prompts several times offers better control and efficiency.) | +| ignore_eos | `bool = False` | Don't stop generation when EOS token is sampled. | +| skip_special_tokens | `bool = True` | Remove special tokens during decoding. | +| spaces_between_special_tokens | `bool = True` | Whether or not to add spaces between special tokens during detokenization. | +| no_stop_trim | `bool = False` | Don't trim stop words or EOS token from the generated text. | +| custom_params | `Optional[List[Optional[Dict[str, Any]]]] = None` | Used when employing `CustomLogitProcessor`. For usage, see below. | + +## Examples + +### Normal + +Launch a server: + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --port 30000 +``` + +Send a request: + +```python +import requests + +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "The capital of France is", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + }, +) +print(response.json()) +``` + +Detailed example in [send request](./send_request.ipynb). + +### Streaming + +Send a request and stream the output: + +```python +import requests, json + +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "The capital of France is", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + "stream": True, + }, + stream=True, +) + +prev = 0 +for chunk in response.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + output = data["text"].strip() + print(output[prev:], end="", flush=True) + prev = len(output) +print("") +``` + +Detailed example in [openai compatible api](openai_api_completions.ipynb). + +### Multimodal + +Launch a server: + +```bash +python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov +``` + +Download an image: + +```bash +curl -o example_image.png -L https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true +``` + +Send a request: + +```python +import requests + +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n\nDescribe this image in a very short sentence.<|im_end|>\n" + "<|im_start|>assistant\n", + "image_data": "example_image.png", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + }, +) +print(response.json()) +``` + +The `image_data` can be a file name, a URL, or a base64 encoded string. See also `python/sglang/srt/utils.py:load_image`. + +Streaming is supported in a similar manner as [above](#streaming). + +Detailed example in [OpenAI API Vision](openai_api_vision.ipynb). + +### Structured Outputs (JSON, Regex, EBNF) + +You can specify a JSON schema, regular expression or [EBNF](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form) to constrain the model output. The model output will be guaranteed to follow the given constraints. Only one constraint parameter (`json_schema`, `regex`, or `ebnf`) can be specified for a request. + +SGLang supports two grammar backends: + +- [XGrammar](https://github.com/mlc-ai/xgrammar) (default): Supports JSON schema, regular expression, and EBNF constraints. + - XGrammar currently uses the [GGML BNF format](https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md). +- [Outlines](https://github.com/dottxt-ai/outlines): Supports JSON schema and regular expression constraints. + +If instead you want to initialize the Outlines backend, you can use `--grammar-backend outlines` flag: + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ +--port 30000 --host 0.0.0.0 --grammar-backend [xgrammar|outlines] # xgrammar or outlines (default: xgrammar) +``` + +```python +import json +import requests + +json_schema = json.dumps({ + "type": "object", + "properties": { + "name": {"type": "string", "pattern": "^[\\w]+$"}, + "population": {"type": "integer"}, + }, + "required": ["name", "population"], +}) + +# JSON (works with both Outlines and XGrammar) +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "Here is the information of the capital of France in the JSON format.\n", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 64, + "json_schema": json_schema, + }, + }, +) +print(response.json()) + +# Regular expression (Outlines backend only) +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "Paris is the capital of", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 64, + "regex": "(France|England)", + }, + }, +) +print(response.json()) + +# EBNF (XGrammar backend only) +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "Write a greeting.", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 64, + "ebnf": 'root ::= "Hello" | "Hi" | "Hey"', + }, + }, +) +print(response.json()) +``` + +Detailed example in [structured outputs](../advanced_features/structured_outputs.ipynb). + +### Custom logit processor + +Launch a server with `--enable-custom-logit-processor` flag on. + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --port 30000 --enable-custom-logit-processor +``` + +Define a custom logit processor that will always sample a specific token id. + +```python +from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor + +class DeterministicLogitProcessor(CustomLogitProcessor): + """A dummy logit processor that changes the logits to always + sample the given token id. + """ + + def __call__(self, logits, custom_param_list): + # Check that the number of logits matches the number of custom parameters + assert logits.shape[0] == len(custom_param_list) + key = "token_id" + + for i, param_dict in enumerate(custom_param_list): + # Mask all other tokens + logits[i, :] = -float("inf") + # Assign highest probability to the specified token + logits[i, param_dict[key]] = 0.0 + return logits +``` + +Send a request: + +```python +import requests + +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "The capital of France is", + "custom_logit_processor": DeterministicLogitProcessor().to_str(), + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": 32, + "custom_params": {"token_id": 5}, + }, + }, +) +print(response.json()) +``` diff --git a/docs/basic_usage/send_request.ipynb b/docs/basic_usage/send_request.ipynb new file mode 100644 index 000000000..6e457a02b --- /dev/null +++ b/docs/basic_usage/send_request.ipynb @@ -0,0 +1,253 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Sending Requests\n", + "This notebook provides a quick-start guide to use SGLang in chat completions after installation.\n", + "\n", + "- For Vision Language Models, see [OpenAI APIs - Vision](openai_api_vision.ipynb).\n", + "- For Embedding Models, see [OpenAI APIs - Embedding](openai_api_embeddings.ipynb) and [Encode (embedding model)](native_api.html#Encode-(embedding-model)).\n", + "- For Reward Models, see [Classify (reward model)](native_api.html#Classify-(reward-model))." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "# This is equivalent to running the following command in your terminal\n", + "# python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct --host 0.0.0.0\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct \\\n", + " --host 0.0.0.0 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using cURL\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess, json\n", + "\n", + "curl_command = f\"\"\"\n", + "curl -s http://localhost:{port}/v1/chat/completions \\\n", + " -H \"Content-Type: application/json\" \\\n", + " -d '{{\"model\": \"qwen/qwen2.5-0.5b-instruct\", \"messages\": [{{\"role\": \"user\", \"content\": \"What is the capital of France?\"}}]}}'\n", + "\"\"\"\n", + "\n", + "response = json.loads(subprocess.check_output(curl_command, shell=True))\n", + "print_highlight(response)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Requests" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "url = f\"http://localhost:{port}/v1/chat/completions\"\n", + "\n", + "data = {\n", + " \"model\": \"qwen/qwen2.5-0.5b-instruct\",\n", + " \"messages\": [{\"role\": \"user\", \"content\": \"What is the capital of France?\"}],\n", + "}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using OpenAI Python Client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"qwen/qwen2.5-0.5b-instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "print_highlight(response)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Streaming" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "# Use stream=True for streaming responses\n", + "response = client.chat.completions.create(\n", + " model=\"qwen/qwen2.5-0.5b-instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + " stream=True,\n", + ")\n", + "\n", + "# Handle the streaming output\n", + "for chunk in response:\n", + " if chunk.choices[0].delta.content:\n", + " print(chunk.choices[0].delta.content, end=\"\", flush=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Native Generation APIs\n", + "\n", + "You can also use the native `/generate` endpoint with requests, which provides more flexibility. An API reference is available at [Sampling Parameters](sampling_params.md)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "response = requests.post(\n", + " f\"http://localhost:{port}/generate\",\n", + " json={\n", + " \"text\": \"The capital of France is\",\n", + " \"sampling_params\": {\n", + " \"temperature\": 0,\n", + " \"max_new_tokens\": 32,\n", + " },\n", + " },\n", + ")\n", + "\n", + "print_highlight(response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Streaming" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests, json\n", + "\n", + "response = requests.post(\n", + " f\"http://localhost:{port}/generate\",\n", + " json={\n", + " \"text\": \"The capital of France is\",\n", + " \"sampling_params\": {\n", + " \"temperature\": 0,\n", + " \"max_new_tokens\": 32,\n", + " },\n", + " \"stream\": True,\n", + " },\n", + " stream=True,\n", + ")\n", + "\n", + "prev = 0\n", + "for chunk in response.iter_lines(decode_unicode=False):\n", + " chunk = chunk.decode(\"utf-8\")\n", + " if chunk and chunk.startswith(\"data:\"):\n", + " if chunk == \"data: [DONE]\":\n", + " break\n", + " data = json.loads(chunk[5:].strip(\"\\n\"))\n", + " output = data[\"text\"]\n", + " print(output[prev:], end=\"\", flush=True)\n", + " prev = len(output)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 000000000..d6ca64d88 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,205 @@ +import os +import sys +from datetime import datetime + +sys.path.insert(0, os.path.abspath("../..")) + +version_file = "../python/sglang/version.py" +with open(version_file, "r") as f: + exec(compile(f.read(), version_file, "exec")) +__version__ = locals()["__version__"] + +project = "SGLang" +copyright = f"2023-{datetime.now().year}, SGLang" +author = "SGLang Team" + +version = __version__ +release = __version__ + +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.napoleon", + "sphinx.ext.viewcode", + "sphinx.ext.autosectionlabel", + "sphinx.ext.intersphinx", + "sphinx_tabs.tabs", + "myst_parser", + "sphinx_copybutton", + "sphinxcontrib.mermaid", + "nbsphinx", + "sphinx.ext.mathjax", +] + +nbsphinx_allow_errors = True +nbsphinx_execute = "never" + +autosectionlabel_prefix_document = True +nbsphinx_allow_directives = True + + +myst_enable_extensions = [ + "dollarmath", + "amsmath", + "deflist", + "colon_fence", + "html_image", + "linkify", + "substitution", +] + +myst_heading_anchors = 3 + +nbsphinx_kernel_name = "python3" +nbsphinx_execute_arguments = [ + "--InlineBackend.figure_formats={'svg', 'pdf'}", + "--InlineBackend.rc={'figure.dpi': 96}", +] + + +nb_render_priority = { + "html": ( + "application/vnd.jupyter.widget-view+json", + "application/javascript", + "text/html", + "image/svg+xml", + "image/png", + "image/jpeg", + "text/markdown", + "text/latex", + "text/plain", + ) +} + +myst_enable_extensions = [ + "dollarmath", + "amsmath", + "deflist", + "colon_fence", + "html_image", + "linkify", + "substitution", +] + +myst_heading_anchors = 3 +myst_ref_domains = ["std", "py"] + +templates_path = ["_templates"] + +source_suffix = { + ".rst": "restructuredtext", + ".md": "markdown", +} + +master_doc = "index" + +language = "en" + +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] + +pygments_style = "sphinx" + +html_theme = "sphinx_book_theme" +html_logo = "_static/image/logo.png" +html_favicon = "_static/image/logo.ico" +html_title = project +html_copy_source = True +html_last_updated_fmt = "" + +html_theme_options = { + "repository_url": "https://github.com/sgl-project/sgl-project.github.io", + "repository_branch": "main", + "show_navbar_depth": 3, + "max_navbar_depth": 4, + "collapse_navbar": True, + "use_edit_page_button": True, + "use_source_button": True, + "use_issues_button": True, + "use_repository_button": True, + "use_download_button": True, + "use_sidenotes": True, + "show_toc_level": 2, +} + +html_context = { + "display_github": True, + "github_user": "sgl-project", + "github_repo": "sgl-project.github.io", + "github_version": "main", + "conf_py_path": "/docs/", +} + +html_static_path = ["_static"] +html_css_files = ["css/custom_log.css"] + + +def setup(app): + app.add_css_file("css/custom_log.css") + + +myst_enable_extensions = [ + "dollarmath", + "amsmath", + "deflist", + "colon_fence", +] +myst_heading_anchors = 5 + +htmlhelp_basename = "sglangdoc" + +latex_elements = {} + +latex_documents = [ + (master_doc, "sglang.tex", "sglang Documentation", "SGLang Team", "manual"), +] + +man_pages = [(master_doc, "sglang", "sglang Documentation", [author], 1)] + +texinfo_documents = [ + ( + master_doc, + "sglang", + "sglang Documentation", + author, + "sglang", + "One line description of project.", + "Miscellaneous", + ), +] + +epub_title = project + +epub_exclude_files = ["search.html"] + +copybutton_prompt_text = r">>> |\.\.\. " +copybutton_prompt_is_regexp = True + +autodoc_preserve_defaults = True +navigation_with_keys = False + +autodoc_mock_imports = [ + "torch", + "transformers", + "triton", +] + +intersphinx_mapping = { + "python": ("https://docs.python.org/3.12", None), + "typing_extensions": ("https://typing-extensions.readthedocs.io/en/latest", None), + "pillow": ("https://pillow.readthedocs.io/en/stable", None), + "numpy": ("https://numpy.org/doc/stable", None), + "torch": ("https://pytorch.org/docs/stable", None), +} + +html_theme = "sphinx_book_theme" + + +nbsphinx_prolog = """ +.. raw:: html + + +""" diff --git a/docs/deploy.py b/docs/deploy.py new file mode 100644 index 000000000..75b7ea7f2 --- /dev/null +++ b/docs/deploy.py @@ -0,0 +1,22 @@ +# Deploy the documents + +import os +from datetime import datetime + + +def run_cmd(cmd): + print(cmd) + os.system(cmd) + + +run_cmd("cd $DOC_SITE_PATH; git pull") + +# (Optional) Remove old files +# run_cmd("rm -rf $ALPA_SITE_PATH/*") + +run_cmd("cp -r _build/html/* $DOC_SITE_PATH") + +cmd_message = f"Update {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" +run_cmd( + f"cd $DOC_SITE_PATH; git add .; git commit -m '{cmd_message}'; git push origin main" +) diff --git a/docs/developer_guide/bench_serving.md b/docs/developer_guide/bench_serving.md new file mode 100644 index 000000000..82f7aa2af --- /dev/null +++ b/docs/developer_guide/bench_serving.md @@ -0,0 +1,334 @@ +## Bench Serving Guide + +This guide explains how to benchmark online serving throughput and latency using `python -m sglang.bench_serving`. It supports multiple inference backends via OpenAI-compatible and native endpoints, and produces both console metrics and optional JSONL outputs. + +### What it does + +- Generates synthetic or dataset-driven prompts and submits them to a target serving endpoint +- Measures throughput, time-to-first-token (TTFT), inter-token latency (ITL), per-request end-to-end latency, and more +- Supports streaming or non-streaming modes, rate control, and concurrency limits + +### Supported backends and endpoints + +- `sglang` / `sglang-native`: `POST /generate` +- `sglang-oai`, `vllm`, `lmdeploy`: `POST /v1/completions` +- `sglang-oai-chat`, `vllm-chat`, `lmdeploy-chat`: `POST /v1/chat/completions` +- `trt` (TensorRT-LLM): `POST /v2/models/ensemble/generate_stream` +- `gserver`: Custom server (Not Implemented yet in this script) +- `truss`: `POST /v1/models/model:predict` + +If `--base-url` is provided, requests are sent to it. Otherwise, `--host` and `--port` are used. When `--model` is not provided, the script will attempt to query `GET /v1/models` for an available model ID (OpenAI-compatible endpoints). + +### Prerequisites + +- Python 3.8+ +- Dependencies typically used by this script: `aiohttp`, `numpy`, `requests`, `tqdm`, `transformers`, and for some datasets `datasets`, `pillow`, `pybase64`. Install as needed. +- An inference server running and reachable via the endpoints above +- If your server requires authentication, set environment variable `OPENAI_API_KEY` (used as `Authorization: Bearer `) + +### Quick start + +Run a basic benchmark against an sglang server exposing `/generate`: + +```bash +python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct +``` + +```bash +python3 -m sglang.bench_serving \ + --backend sglang \ + --host 127.0.0.1 --port 30000 \ + --num-prompts 1000 \ + --model meta-llama/Llama-3.1-8B-Instruct +``` + +Or, using an OpenAI-compatible endpoint (completions): + +```bash +python3 -m sglang.bench_serving \ + --backend vllm \ + --base-url http://127.0.0.1:8000 \ + --num-prompts 1000 \ + --model meta-llama/Llama-3.1-8B-Instruct +``` + +### Datasets + +Select with `--dataset-name`: + +- `sharegpt` (default): loads ShareGPT-style pairs; optionally restrict with `--sharegpt-context-len` and override outputs with `--sharegpt-output-len` +- `random`: random text lengths; sampled from ShareGPT token space +- `random-ids`: random token ids (can lead to gibberish) +- `random-image`: generates random images and wraps them in chat messages; supports custom resolutions via 'heightxwidth' format +- `generated-shared-prefix`: synthetic dataset with shared long system prompts and short questions +- `mmmu`: samples from MMMU (Math split) and includes images + +Common dataset flags: + +- `--num-prompts N`: number of requests +- `--random-input-len`, `--random-output-len`, `--random-range-ratio`: for random/random-ids/random-image +- `--random-image-num-images`, `--random-image-resolution`: for random-image dataset (supports presets 1080p/720p/360p or custom 'heightxwidth' format) +- `--apply-chat-template`: apply tokenizer chat template when constructing prompts +- `--dataset-path PATH`: file path for ShareGPT json; if blank and missing, it will be downloaded and cached + +Generated Shared Prefix flags (for `generated-shared-prefix`): + +- `--gsp-num-groups` +- `--gsp-prompts-per-group` +- `--gsp-system-prompt-len` +- `--gsp-question-len` +- `--gsp-output-len` + +Random Image dataset flags (for `random-image`): + +- `--random-image-num-images`: Number of images per request +- `--random-image-resolution`: Image resolution; supports presets (1080p, 720p, 360p) or custom 'heightxwidth' format (e.g., 1080x1920, 512x768) + +### Examples + +1. To benchmark random-image dataset with 3 images per request, 500 prompts, 512 input length, and 512 output length, you can run: + +```bash +python -m sglang.launch_server --model-path Qwen/Qwen2.5-VL-3B-Instruct --disable-radix-cache +``` + +```bash +python -m sglang.bench_serving \ + --backend sglang-oai-chat \ + --dataset-name random-image \ + --num-prompts 500 \ + --random-image-num-images 3 \ + --random-image-resolution 720p \ + --random-input-len 512 \ + --random-output-len 512 +``` + +2. To benchmark random dataset with 3000 prompts, 1024 input length, and 1024 output length, you can run: + +```bash +python -m sglang.launch_server --model-path Qwen/Qwen2.5-3B-Instruct +``` + +```bash +python3 -m sglang.bench_serving \ + --backend sglang \ + --dataset-name random \ + --num-prompts 3000 \ + --random-input 1024 \ + --random-output 1024 \ + --random-range-ratio 0.5 +``` + +### Choosing model and tokenizer + +- `--model` is required unless the backend exposes `GET /v1/models`, in which case the first model ID is auto-selected. +- `--tokenizer` defaults to `--model`. Both can be HF model IDs or local paths. +- For ModelScope workflows, setting `SGLANG_USE_MODELSCOPE=true` enables fetching via ModelScope (weights are skipped for speed). +- If your tokenizer lacks a chat template, the script warns because token counting can be less robust for gibberish outputs. + +### Rate, concurrency, and streaming + +- `--request-rate`: requests per second. `inf` sends all immediately (burst). Non-infinite rate uses a Poisson process for arrival times. +- `--max-concurrency`: caps concurrent in-flight requests regardless of arrival rate. +- `--disable-stream`: switch to non-streaming mode when supported; TTFT then equals total latency for chat completions. + +### Other key options + +- `--output-file FILE.jsonl`: append JSONL results to file; auto-named if unspecified +- `--output-details`: include per-request arrays (generated texts, errors, ttfts, itls, input/output lens) +- `--extra-request-body '{"top_p":0.9,"temperature":0.6}'`: merged into payload (sampling params, etc.) +- `--disable-ignore-eos`: pass through EOS behavior (varies by backend) +- `--warmup-requests N`: run warmup requests with short output first (default 1) +- `--flush-cache`: call `/flush_cache` (sglang) before main run +- `--profile`: call `/start_profile` and `/stop_profile` (requires server to enable profiling, e.g., `SGLANG_TORCH_PROFILER_DIR`) +- `--lora-name name1 name2 ...`: randomly pick one per request and pass to backend (e.g., `lora_path` for sglang) +- `--tokenize-prompt`: send integer IDs instead of text (currently supports `--backend sglang` only) + +### Authentication + +If your target endpoint requires OpenAI-style auth, set: + +```bash +export OPENAI_API_KEY=sk-...yourkey... +``` + +The script will add `Authorization: Bearer $OPENAI_API_KEY` automatically for OpenAI-compatible routes. + +### Metrics explained + +Printed after each run: + +- Request throughput (req/s) +- Input token throughput (tok/s) +- Output token throughput (tok/s) +- Total token throughput (tok/s) +- Concurrency: aggregate time of all requests divided by wall time +- End-to-End Latency (ms): mean/median/std/p99 per-request total latency +- Time to First Token (TTFT, ms): mean/median/std/p99 for streaming mode +- Inter-Token Latency (ITL, ms): mean/median/std/p95/p99/max between tokens +- TPOT (ms): Token processing time after first token, i.e., `(latency - ttft)/(tokens-1)` +- Accept length (sglang-only, if available): speculative decoding accept length + +The script also retokenizes generated text with the configured tokenizer and reports "retokenized" counts. + +### JSONL output format + +When `--output-file` is set, one JSON object is appended per run. Base fields: + +- Arguments summary: backend, dataset, request_rate, max_concurrency, etc. +- Duration and totals: completed, total_input_tokens, total_output_tokens, retokenized totals +- Throughputs and latency statistics as printed in the console +- `accept_length` when available (sglang) + +With `--output-details`, an extended object also includes arrays: + +- `input_lens`, `output_lens` +- `ttfts`, `itls` (per request: ITL arrays) +- `generated_texts`, `errors` + +### End-to-end examples + +1) sglang native `/generate` (streaming): + +```bash +python3 -m sglang.bench_serving \ + --backend sglang \ + --host 127.0.0.1 --port 30000 \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --dataset-name random \ + --random-input-len 1024 --random-output-len 1024 --random-range-ratio 0.5 \ + --num-prompts 2000 \ + --request-rate 100 \ + --max-concurrency 512 \ + --output-file sglang_random.jsonl --output-details +``` + +2) OpenAI-compatible Completions (e.g., vLLM): + +```bash +python3 -m sglang.bench_serving \ + --backend vllm \ + --base-url http://127.0.0.1:8000 \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --dataset-name sharegpt \ + --num-prompts 1000 \ + --sharegpt-output-len 256 +``` + +3) OpenAI-compatible Chat Completions (streaming): + +```bash +python3 -m sglang.bench_serving \ + --backend vllm-chat \ + --base-url http://127.0.0.1:8000 \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --dataset-name random \ + --num-prompts 500 \ + --apply-chat-template +``` + +4) Random images (VLM) with chat template: + +```bash +python3 -m sglang.bench_serving \ + --backend sglang \ + --host 127.0.0.1 --port 30000 \ + --model your-vlm-model \ + --dataset-name random-image \ + --random-image-num-images 2 \ + --random-image-resolution 720p \ + --random-input-len 128 --random-output-len 256 \ + --num-prompts 200 \ + --apply-chat-template +``` + +4a) Random images with custom resolution: + +```bash +python3 -m sglang.bench_serving \ + --backend sglang \ + --host 127.0.0.1 --port 30000 \ + --model your-vlm-model \ + --dataset-name random-image \ + --random-image-num-images 1 \ + --random-image-resolution 512x768 \ + --random-input-len 64 --random-output-len 128 \ + --num-prompts 100 \ + --apply-chat-template +``` + +5) Generated shared prefix (long system prompts + short questions): + +```bash +python3 -m sglang.bench_serving \ + --backend sglang \ + --host 127.0.0.1 --port 30000 \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --dataset-name generated-shared-prefix \ + --gsp-num-groups 64 --gsp-prompts-per-group 16 \ + --gsp-system-prompt-len 2048 --gsp-question-len 128 --gsp-output-len 256 \ + --num-prompts 1024 +``` + +6) Tokenized prompts (ids) for strict length control (sglang only): + +```bash +python3 -m sglang.bench_serving \ + --backend sglang \ + --host 127.0.0.1 --port 30000 \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --dataset-name random \ + --tokenize-prompt \ + --random-input-len 2048 --random-output-len 256 --random-range-ratio 0.2 +``` + +7) Profiling and cache flush (sglang): + +```bash +python3 -m sglang.bench_serving \ + --backend sglang \ + --host 127.0.0.1 --port 30000 \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --profile \ + --flush-cache +``` + +8) TensorRT-LLM streaming endpoint: + +```bash +python3 -m sglang.bench_serving \ + --backend trt \ + --base-url http://127.0.0.1:8000 \ + --model your-trt-llm-model \ + --dataset-name random \ + --num-prompts 100 \ + --disable-ignore-eos +``` + +9) Evaluating large-scale KVCache sharing with mooncake trace (sglang only): + +```bash +python3 -m sglang.bench_serving \ + --backend sglang \ + --host 127.0.0.1 --port 30000 \ + --model mode-name \ + --dataset-name mooncake \ + --mooncake-slowdown-factor 1.0 \ + --mooncake-num-rounds 1000 \ + --mooncake-workload conversation|mooncake|agent|synthetic + --use-trace-timestamps true \ + --random-output-len 256 +``` + +### Troubleshooting + +- All requests failed: verify `--backend`, server URL/port, `--model`, and authentication. Check warmup errors printed by the script. +- Throughput seems too low: adjust `--request-rate` and `--max-concurrency`; verify server batch size/scheduling; ensure streaming is enabled if appropriate. +- Token counts look odd: prefer chat/instruct models with proper chat templates; otherwise tokenization of gibberish may be inconsistent. +- Random-image/MMMU datasets: ensure you installed extra deps (`pillow`, `datasets`, `pybase64`). +- Authentication errors (401/403): set `OPENAI_API_KEY` or disable auth on your server. + +### Notes + +- The script raises the file descriptor soft limit (`RLIMIT_NOFILE`) to help with many concurrent connections. +- For sglang, `/get_server_info` is queried post-run to report speculative decoding accept length when available. diff --git a/docs/developer_guide/benchmark_and_profiling.md b/docs/developer_guide/benchmark_and_profiling.md new file mode 100644 index 000000000..948c837ff --- /dev/null +++ b/docs/developer_guide/benchmark_and_profiling.md @@ -0,0 +1,182 @@ +# Benchmark and Profiling + +## Benchmark + +- Benchmark the latency of running a single static batch without a server. The arguments are the same as for `launch_server.py`. + Note that this is a simplified test script without a dynamic batching server, so it may run out of memory for a batch size that a real server can handle. A real server truncates the prefill into several batches, while this simplified script does not. + - Without a server (do not need to launch a server) + ```bash + python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch 32 --input-len 256 --output-len 32 + ``` + - With a server (please use `sglang.launch_server` to launch a server first and run the following command.) + ```bash + python -m sglang.bench_one_batch_server --base-url http://127.0.0.1:30000 --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch-size 32 --input-len 256 --output-len 32 + ``` + + +- Benchmark offline processing. This script will start an offline engine and run the benchmark. + + ```bash + python3 -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10 + ``` + +- Benchmark online serving. Please use `sglang.launch_server` to launch a server first and run the following command. + + ```bash + python3 -m sglang.bench_serving --backend sglang --num-prompt 10 + ``` + +## Profile with PyTorch Profiler + +[Pytorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) is a convenient basic tool to inspect kernel execution time, call stack, and kernel overlap and occupancy. + +### Profile a server with `sglang.bench_serving` + +```bash +# set trace path +export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log + +# start server +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct + +# send profiling request from client +python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 10 --sharegpt-output-len 100 --profile +``` + +Please make sure that the `SGLANG_TORCH_PROFILER_DIR` should be set at both server and client side, otherwise the trace file cannot be generated correctly . A secure way will be setting `SGLANG_TORCH_PROFILER_DIR` in the `.*rc` file of shell (e.g. `~/.bashrc` for bash shells). + +For more details, please refer to [Bench Serving Guide](./bench_serving.md). + +### Profile a server with `sglang.bench_offline_throughput` +```bash +export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log + +# profile one batch with bench_one_batch.py +# batch size can be controlled with --batch argument +python3 -m sglang.bench_one_batch --model-path meta-llama/Llama-3.1-8B-Instruct --batch 32 --input-len 1024 --output-len 10 --profile + +# profile multiple batches with bench_offline_throughput.py +python -m sglang.bench_offline_throughput --model-path meta-llama/Llama-3.1-8B-Instruct --dataset-name random --num-prompts 10 --profile --mem-frac=0.8 +``` + +### Profile a server with `sglang.profiler` + +When the server is running (e.g., processing a decoding request), you can start live profiling immediately by sending a profile request to the server. + +You can do this by running `python3 -m sglang.profiler`. For example: + +``` +# Terminal 1: Send a generation request +python3 -m sglang.test.send_one + +# Terminal 2: Before the above request finishes, quickly launch the following command in a separate terminal. +# It will generate a profile of the above request for several decoding batches. +python3 -m sglang.profiler +``` + +### Possible PyTorch bugs +If in any cases you encounter the following error (for example, using qwen 2.5 VL): +```bash +RuntimeError: !stack.empty() INTERNAL ASSERT FAILED at "/pytorch/torch/csrc/autograd/profiler_python.cpp":983, please report a bug to PyTorch. Python replay stack is empty. +``` +This is likely a PyTorch Bug reported in [Bug: vLLM Profiler](https://github.com/vllm-project/vllm/issues/18240) and [Bug: torch.profiler.profile](https://github.com/pytorch/pytorch/issues/101632). As a workaround, you may disable `with_stack` with an environment variable such as follows: +```bash +export SGLANG_PROFILE_WITH_STACK=False +python -m sglang.bench_offline_throughput --model-path meta-llama/Llama-3.1-8B-Instruct --dataset-name random --num-prompts 10 --profile --mem-frac=0.8 +``` + +### View traces + +Trace files can be loaded and visualized from: + +1. https://ui.perfetto.dev/ (any browser) +2. chrome://tracing (Chrome browser only) + +If browser cannot open trace file due to its large size, +client can generate a small trace file (<100MB) by controlling number of prompts and lengths of prompt outputs. +For example, when profiling a server, + +```bash +python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 2 --sharegpt-output-len 100 --profile +``` + +This command sets the number of prompts to 2 with `--num-prompts` argument and limits the length of output sequences to 100 with `--sharegpt-output-len` argument, which can generate a small trace file for browser to open smoothly. + +Additionally, if you want to locate the SGLang Python source code through the cuda kernel in Trace, you need to disable CUDA Graph when starting the service. This can be done by using the `--disable-cuda-graph` parameter in the command to start the service. + +## Profile with Nsight + +[Nsight systems](https://docs.nvidia.com/nsight-systems/) is an advanced tool that exposes more profiling details, such as register and shared memory usage, annotated code regions and low-level CUDA APIs and events. + +1. Prerequisite: + + Install using apt, or run inside a [NVIDIA Docker container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags) or [SGLang Docker container](https://github.com/sgl-project/sglang/tree/main/docker). + + ```bash + # install nsys + # https://docs.nvidia.com/nsight-systems/InstallationGuide/index.html + apt update + apt install -y --no-install-recommends gnupg + echo "deb http://developer.download.nvidia.com/devtools/repos/ubuntu$(source /etc/lsb-release; echo "$DISTRIB_RELEASE" | tr -d .)/$(dpkg --print-architecture) /" | tee /etc/apt/sources.list.d/nvidia-devtools.list + apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub + apt update + apt install nsight-systems-cli + ``` + +2. To profile a single batch, use + + ```bash + nsys profile --trace-fork-before-exec=true --cuda-graph-trace=node python3 -m sglang.bench_one_batch --model meta-llama/Meta-Llama-3-8B --batch-size 64 --input-len 512 + ``` + +3. To profile a server, e.g. + + ```bash + # launch the server, set the delay and duration times according to needs + # after the duration time has been used up, server will be killed by nsys + + nsys profile --trace-fork-before-exec=true --cuda-graph-trace=node -o sglang.out --delay 60 --duration 70 python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disable-radix-cache + + # client + python3 -m sglang.bench_serving --backend sglang --num-prompts 1000 --dataset-name random --random-input 1024 --random-output 512 + ``` + + In practice, we recommend users to set `--duration` argument to a large value. Whenever user wants the server to stop profiling. Firstly run: + + ```bash + nsys sessions list + ``` + + to get the session id in the form of `profile-XXXXX`, then run: + + ```bash + nsys stop --session=profile-XXXXX + ``` + + to manually kill the profiler and generate `nsys-rep` files instantly. + +4. Use NVTX to annotate code regions, e.g. to see their execution time. + + ```bash + # install nvtx + pip install nvtx + ``` + + ```python + # code snippets + import nvtx + with nvtx.annotate("description", color="color"): + # some critical code + ``` + +## Other tips + +1. You can benchmark a model using dummy weights by only providing the config.json file. This allows for quick testing of model variants without training. To do so, add `--load-format dummy` to the above commands and then you only need a correct `config.json` under the checkpoint folder. +2. You can benchmark a model with modified configs (e.g., less layers) by using `--json-model-override-args`. For example, you can benchmark a model with only 2 layers and 2 kv heads using: + + ```bash + python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch 32 --input-len 256 --output-len 32 --load-format dummy --json-model-override-args '{"num_hidden_layers": 1, "num_key_value_heads": 1}' + ``` + +3. You can use `--python-backtrace=cuda` to see python call stack for all CUDA kernels, as in PyTorch Profiler. (Caveat: this can cause inaccurately long kernel runtimes for CUDA event based timing) +4. For more arguments see [Nsight Systems User Guide](https://docs.nvidia.com/nsight-systems/UserGuide/index.html). diff --git a/docs/developer_guide/contribution_guide.md b/docs/developer_guide/contribution_guide.md new file mode 100644 index 000000000..f8e6f692d --- /dev/null +++ b/docs/developer_guide/contribution_guide.md @@ -0,0 +1,103 @@ +# Contribution Guide + +Welcome to **SGLang**! We appreciate your interest in contributing. This guide provides a concise overview of how to set up your environment, run tests, build documentation, and open a Pull Request (PR). Whether you’re fixing a small bug or developing a major feature, we encourage following these steps for a smooth contribution process. + +## Install SGLang from Source + +### Fork and clone the repository + +**Note**: New contributors do **not** have the write permission to push to the official SGLang repo. Please fork the repository under your GitHub account, then clone your fork locally. + +```bash +git clone https://github.com//sglang.git +``` + +### Build from source + +Refer to [Install SGLang from Source](../get_started/install.md#method-2-from-source). + +## Format code with pre-commit + +We use [pre-commit](https://pre-commit.com/) to maintain consistent code style checks. Before pushing your changes, please run: + +```bash +pip3 install pre-commit +pre-commit install +pre-commit run --all-files +``` + +- **`pre-commit run --all-files`** manually runs all configured checks, applying fixes if possible. If it fails the first time, re-run it to ensure lint errors are fully resolved. Make sure your code passes all checks **before** creating a Pull Request. +- **Do not commit** directly to the `main` branch. Always create a new branch (e.g., `feature/my-new-feature`), push your changes, and open a PR from that branch. + +## Run and add unit tests + +If you add a new feature or fix a bug, please add corresponding unit tests to ensure coverage and prevent regression. +SGLang uses Python's built-in [unittest](https://docs.python.org/3/library/unittest.html) framework. +For detailed instructions on running tests and integrating them into CI, refer to [test/README.md](https://github.com/sgl-project/sglang/tree/main/test/README.md). + +## Write documentations + +We recommend new contributors start from writing documentation, which helps you quickly understand SGLang codebase. +For more details, please refer to [docs/README.md](https://github.com/sgl-project/sglang/tree/main/docs/README.md). + +## Test the accuracy +If your code changes the model output, please run the accuracy tests. A quick sanity check is the few-shot GSM8K. + +``` +# Launch a server +python3 -m sglang.launch_server --model Qwen/Qwen2-7B-Instruct + +# Evaluate +python3 -m sglang.test.few_shot_gsm8k --num-questions 200 +``` + +Please note that the above script is primarily a sanity check, not a rigorous accuracy or speed test. +This test can have significant variance (1%–5%) in accuracy due to batching and the non-deterministic nature of the inference engine. +Also, do not rely on the "Latency/Output throughput" from this script, as it is not a proper speed test. + +GSM8K is too easy for state-of-the-art models nowadays. Please try your own more challenging accuracy tests. +You can find additional accuracy eval examples in: +- [test_eval_accuracy_large.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_eval_accuracy_large.py) +- [test_gpt_oss_1gpu.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_gpt_oss_1gpu.py) + +## Benchmark the speed +Refer to [Benchmark and Profiling](../developer_guide/benchmark_and_profiling.md). + +## Request a review +You can identify potential reviewers for your code by checking the [code owners](https://github.com/sgl-project/sglang/blob/main/.github/CODEOWNERS) and [reviewers](https://github.com/sgl-project/sglang/blob/main/.github/REVIEWERS.md) files. +Another effective strategy is to review the file modification history and contact individuals who have frequently edited the files. +If you modify files protected by code owners, their approval is required to merge the code. + +## General code style +- Avoid code duplication. If the same code snippet (more than five lines) appears multiple times, extract it into a shared function. +- Minimize device synchronization. Reduce expensive CPU-GPU synchronization operations, such as `tensor.item()` or `tensor.cpu()`, whenever possible. Use vectorized code. +- Keep files concise. If a file exceeds 2,000 lines of code, split it into multiple smaller files. +- Prioritize extreme efficiency. SGLang is a runtime, and most of your code runs on the critical path for every request. Optimize all minor overheads as much as possible, especially in the model forward code. + - A common pattern is some runtime checks in the model forward pass (e.g., [this](https://github.com/sgl-project/sglang/blob/f1b0eda55c2c4838e8ab90a0fac7fb1e3d7064ab/python/sglang/srt/models/deepseek_v2.py#L486-L491)). These are very likely the same for every layer. Please cache the result as a single boolean value whenever possible. +- Strive to make functions as pure as possible. Avoid in-place modification of arguments. +- When supporting new hardware or features, follow these guidelines: + - Do not drastically change existing code. + - Always prefer new files to introduce specific components for your new hardware (e.g., `allocator_ascend.py`). + - If you write multiple if/else blocks for new features, ensure the common path (e.g., NVIDIA hardware or the existing code path) is the first branch. + +## How to update sgl-kernel +Since sglang and sgl-kernel are separate Python packages, our current GitHub CI infrastructure does not support updating a kernel and using it immediately within the same pull request (PR). +To add a new kernel or modify an existing one in the sgl-kernel package, you must use multiple PRs. + +Follow these steps: + +1. Submit a PR to update the sgl-kernel source code without using it in sglang python package (e.g., [#8884](https://github.com/sgl-project/sglang/pull/8884/files)). +2. Bump the version of sgl-kernel (e.g., [#9220](https://github.com/sgl-project/sglang/pull/9220/files)). + - Once merged, this will trigger an automatic release of the sgl-kernel wheel to PyPI. + - If not urgent, you can wait for other people to release the wheel. A new version will typically be released within one week. +3. Apply the changes: + - Update the sgl-kernel version in `sglang/python/pyproject.toml` to use the modified kernels. + - Update the related caller code in the sglang to use the new kernel. + +## Tips for newcomers + +If you want to contribute but don’t have a specific idea in mind, pick issues labeled [“good first issue” or “help wanted”](https://github.com/sgl-project/sglang/issues?q=is%3Aissue+label%3A%22good+first+issue%22%2C%22help+wanted%22). These tasks typically have lower complexity and provide an excellent introduction to the codebase. Also check out this [code walk-through](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/tree/main/sglang/code-walk-through) for a deeper look into SGLang’s workflow. + +If you have any questions or want to start a discussion, please feel free to ask in our [Slack channel](https://slack.sglang.ai). + +Thank you for your interest in SGLang. Happy coding! diff --git a/docs/developer_guide/development_guide_using_docker.md b/docs/developer_guide/development_guide_using_docker.md new file mode 100644 index 000000000..e38947902 --- /dev/null +++ b/docs/developer_guide/development_guide_using_docker.md @@ -0,0 +1,108 @@ +# Development Guide Using Docker + +## Setup VSCode on a Remote Host +(Optional - you can skip this step if you plan to run sglang dev container locally) + +1. In the remote host, download `code` from [Https://code.visualstudio.com/docs/?dv=linux64cli](https://code.visualstudio.com/download) and run `code tunnel` in a shell. + +Example +```bash +wget https://vscode.download.prss.microsoft.com/dbazure/download/stable/fabdb6a30b49f79a7aba0f2ad9df9b399473380f/vscode_cli_alpine_x64_cli.tar.gz +tar xf vscode_cli_alpine_x64_cli.tar.gz + +# https://code.visualstudio.com/docs/remote/tunnels +./code tunnel +``` + +2. In your local machine, press F1 in VSCode and choose "Remote Tunnels: Connect to Tunnel". + +## Setup Docker Container + +### Option 1. Use the default dev container automatically from VSCode +There is a `.devcontainer` folder in the sglang repository root folder to allow VSCode to automatically start up within dev container. You can read more about this VSCode extension in VSCode official document [Developing inside a Container](https://code.visualstudio.com/docs/devcontainers/containers). +![image](https://github.com/user-attachments/assets/6a245da8-2d4d-4ea8-8db1-5a05b3a66f6d) +(*Figure 1: Diagram from VSCode official documentation [Developing inside a Container](https://code.visualstudio.com/docs/devcontainers/containers).*) + +To enable this, you only need to: +1. Start Visual Studio Code and install [VSCode dev container extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers). +2. Press F1, type and choose "Dev Container: Open Folder in Container. +3. Input the `sglang` local repo path in your machine and press enter. + +The first time you open it in dev container might take longer due to docker pull and build. Once it's successful, you should set on your status bar at the bottom left displaying that you are in a dev container: + +![image](https://github.com/user-attachments/assets/650bba0b-c023-455f-91f9-ab357340106b) + +Now when you run `sglang.launch_server` in the VSCode terminal or start debugging using F5, sglang server will be started in the dev container with all your local changes applied automatically: + +![image](https://github.com/user-attachments/assets/748c85ba-7f8c-465e-8599-2bf7a8dde895) + + +### Option 2. Start up containers manually (advanced) + +The following startup command is an example for internal development by the SGLang team. You can **modify or add directory mappings as needed**, especially for model weight downloads, to prevent repeated downloads by different Docker containers. + +❗️ **Note on RDMA** + + 1. `--network host` and `--privileged` are required by RDMA. If you don't need RDMA, you can remove them but keeping them there does not harm. Thus, we enable these two flags by default in the commands below. + 2. You may need to set `NCCL_IB_GID_INDEX` if you are using RoCE, for example: `export NCCL_IB_GID_INDEX=3`. + +```bash +# Change the name to yours +docker run -itd --shm-size 32g --gpus all -v --ipc=host --network=host --privileged --name sglang_dev lmsysorg/sglang:dev /bin/zsh +docker exec -it sglang_dev /bin/zsh +``` +Some useful volumes to mount are: +1. **Huggingface model cache**: mounting model cache can avoid re-download every time docker restarts. Default location on Linux is `~/.cache/huggingface/`. +2. **SGLang repository**: code changes in the SGLang local repository will be automatically synced to the .devcontainer. + +Example 1: Monting local cache folder `/opt/dlami/nvme/.cache` but not the SGLang repo. Use this when you prefer to manually transfer local code changes to the devcontainer. +```bash +docker run -itd --shm-size 32g --gpus all -v /opt/dlami/nvme/.cache:/root/.cache --ipc=host --network=host --privileged --name sglang_zhyncs lmsysorg/sglang:dev /bin/zsh +docker exec -it sglang_zhyncs /bin/zsh +``` +Example 2: Mounting both HuggingFace cache and local SGLang repo. Local code changes are automatically synced to the devcontainer as the SGLang is installed in editable mode in the dev image. +```bash +docker run -itd --shm-size 32g --gpus all -v $HOME/.cache/huggingface/:/root/.cache/huggingface -v $HOME/src/sglang:/sgl-workspace/sglang --ipc=host --network=host --privileged --name sglang_zhyncs lmsysorg/sglang:dev /bin/zsh +docker exec -it sglang_zhyncs /bin/zsh +``` +## Debug SGLang with VSCode Debugger +1. (Create if not exist) open `launch.json` in VSCode. +2. Add the following config and save. Please note that you can edit the script as needed to apply different parameters or debug a different program (e.g. benchmark script). + ```JSON + { + "version": "0.2.0", + "configurations": [ + { + "name": "Python Debugger: launch_server", + "type": "debugpy", + "request": "launch", + "module": "sglang.launch_server", + "console": "integratedTerminal", + "args": [ + "--model-path", "meta-llama/Llama-3.2-1B", + "--host", "0.0.0.0", + "--port", "30000", + "--trust-remote-code", + ], + "justMyCode": false + } + ] + } + ``` + +3. Press "F5" to start. VSCode debugger will ensure that the program will pause at the breakpoints even if the program is running at remote SSH/Tunnel host + dev container. + +## Profile + +```bash +# Change batch size, input, output and add `disable-cuda-graph` (for easier analysis) +# e.g. DeepSeek V3 +nsys profile -o deepseek_v3 python3 -m sglang.bench_one_batch --batch-size 1 --input 128 --output 256 --model deepseek-ai/DeepSeek-V3 --trust-remote-code --tp 8 --disable-cuda-graph +``` + +## Evaluation + +```bash +# e.g. gsm8k 8 shot +python3 benchmark/gsm8k/bench_sglang.py --num-questions 2000 --parallel 2000 --num-shots 8 +``` diff --git a/docs/developer_guide/release_process.md b/docs/developer_guide/release_process.md new file mode 100644 index 000000000..e817a256e --- /dev/null +++ b/docs/developer_guide/release_process.md @@ -0,0 +1,18 @@ +# PyPI Package Release Process + +## Update the version in code +Update the package version in `python/pyproject.toml` and `python/sglang/__init__.py`. + +## Upload the PyPI package + +``` +pip install build twine +``` + +``` +cd python +bash upload_pypi.sh +``` + +## Make a release in GitHub +Make a new release https://github.com/sgl-project/sglang/releases/new. diff --git a/docs/developer_guide/setup_github_runner.md b/docs/developer_guide/setup_github_runner.md new file mode 100644 index 000000000..6ed78a247 --- /dev/null +++ b/docs/developer_guide/setup_github_runner.md @@ -0,0 +1,49 @@ +# Set Up Self-Hosted Runners for GitHub Action + +## Add a Runner + +### Step 1: Start a docker container. + +You can mount a folder for the shared huggingface model weights cache. The command below uses `/tmp/huggingface` as an example. + +``` +docker pull nvidia/cuda:12.1.1-devel-ubuntu22.04 +# Nvidia +docker run --shm-size 128g -it -v /tmp/huggingface:/hf_home --gpus all nvidia/cuda:12.1.1-devel-ubuntu22.04 /bin/bash +# AMD +docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.5.0rc1-rocm630 /bin/bash +# AMD just the last 2 GPUs +docker run --rm --device=/dev/kfd --device=/dev/dri/renderD176 --device=/dev/dri/renderD184 --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.5.0rc1-rocm630 /bin/bash +``` + +### Step 2: Configure the runner by `config.sh` + +Run these commands inside the container. + +``` +apt update && apt install -y curl python3-pip git +export RUNNER_ALLOW_RUNASROOT=1 +``` + +Then follow https://github.com/sgl-project/sglang/settings/actions/runners/new?arch=x64&os=linux to run `config.sh` + +**Notes** +- Do not need to specify the runner group +- Give it a name (e.g., `test-sgl-gpu-0`) and some labels (e.g., `1-gpu-runner`). The labels can be edited later in Github Settings. +- Do not need to change the work folder. + +### Step 3: Run the runner by `run.sh` + +- Set up environment variables +``` +export HF_HOME=/hf_home +export SGLANG_IS_IN_CI=true +export HF_TOKEN=hf_xxx +export OPENAI_API_KEY=sk-xxx +export CUDA_VISIBLE_DEVICES=0 +``` + +- Run it forever +``` +while true; do ./run.sh; echo "Restarting..."; sleep 2; done +``` diff --git a/docs/get_started/install.md b/docs/get_started/install.md new file mode 100644 index 000000000..e2e780e00 --- /dev/null +++ b/docs/get_started/install.md @@ -0,0 +1,131 @@ +# Install SGLang + +You can install SGLang using one of the methods below. + +This page primarily applies to common NVIDIA GPU platforms. +For other or newer platforms, please refer to the dedicated pages for [NVIDIA Blackwell GPUs](../platforms/blackwell_gpu.md), [AMD GPUs](../platforms/amd_gpu.md), [Intel Xeon CPUs](../platforms/cpu_server.md), [NVIDIA Jetson](../platforms/nvidia_jetson.md), [Ascend NPUs](../platforms/ascend_npu.md). + +## Method 1: With pip or uv + +It is recommended to use uv for faster installation: + +```bash +pip install --upgrade pip +pip install uv +uv pip install "sglang[all]>=0.5.2" +``` + +**Quick fixes to common problems** +- If you encounter `OSError: CUDA_HOME environment variable is not set`. Please set it to your CUDA install root with either of the following solutions: + 1. Use `export CUDA_HOME=/usr/local/cuda-` to set the `CUDA_HOME` environment variable. + 2. Install FlashInfer first following [FlashInfer installation doc](https://docs.flashinfer.ai/installation.html), then install SGLang as described above. + +## Method 2: From source + +```bash +# Use the last release branch +git clone -b v0.5.2 https://github.com/sgl-project/sglang.git +cd sglang + +# Install the python packages +pip install --upgrade pip +pip install -e "python[all]" +``` + +**Quick fixes to common problems** +- If you want to develop SGLang, it is recommended to use docker. Please refer to [setup docker container](../developer_guide/development_guide_using_docker.md#setup-docker-container). The docker image is `lmsysorg/sglang:dev`. + +## Method 3: Using docker + +The docker images are available on Docker Hub at [lmsysorg/sglang](https://hub.docker.com/r/lmsysorg/sglang/tags), built from [Dockerfile](https://github.com/sgl-project/sglang/tree/main/docker). +Replace `` below with your huggingface hub [token](https://huggingface.co/docs/hub/en/security-tokens). + +```bash +docker run --gpus all \ + --shm-size 32g \ + -p 30000:30000 \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --host 0.0.0.0 --port 30000 +``` + +## Method 4: Using Kubernetes + +Please check out [OME](https://github.com/sgl-project/ome), a Kubernetes operator for enterprise-grade management and serving of large language models (LLMs). + +
+More + +1. Option 1: For single node serving (typically when the model size fits into GPUs on one node) + + Execute command `kubectl apply -f docker/k8s-sglang-service.yaml`, to create k8s deployment and service, with llama-31-8b as example. + +2. Option 2: For multi-node serving (usually when a large model requires more than one GPU node, such as `DeepSeek-R1`) + + Modify the LLM model path and arguments as necessary, then execute command `kubectl apply -f docker/k8s-sglang-distributed-sts.yaml`, to create two nodes k8s statefulset and serving service. + +
+ +## Method 5: Using docker compose + +
+More + +> This method is recommended if you plan to serve it as a service. +> A better approach is to use the [k8s-sglang-service.yaml](https://github.com/sgl-project/sglang/blob/main/docker/k8s-sglang-service.yaml). + +1. Copy the [compose.yml](https://github.com/sgl-project/sglang/blob/main/docker/compose.yaml) to your local machine +2. Execute the command `docker compose up -d` in your terminal. +
+ +## Method 6: Run on Kubernetes or Clouds with SkyPilot + +
+More + +To deploy on Kubernetes or 12+ clouds, you can use [SkyPilot](https://github.com/skypilot-org/skypilot). + +1. Install SkyPilot and set up Kubernetes cluster or cloud access: see [SkyPilot's documentation](https://skypilot.readthedocs.io/en/latest/getting-started/installation.html). +2. Deploy on your own infra with a single command and get the HTTP API endpoint: +
+SkyPilot YAML: sglang.yaml + +```yaml +# sglang.yaml +envs: + HF_TOKEN: null + +resources: + image_id: docker:lmsysorg/sglang:latest + accelerators: A100 + ports: 30000 + +run: | + conda deactivate + python3 -m sglang.launch_server \ + --model-path meta-llama/Llama-3.1-8B-Instruct \ + --host 0.0.0.0 \ + --port 30000 +``` + +
+ +```bash +# Deploy on any cloud or Kubernetes cluster. Use --cloud to select a specific cloud provider. +HF_TOKEN= sky launch -c sglang --env HF_TOKEN sglang.yaml + +# Get the HTTP API endpoint +sky status --endpoint 30000 sglang +``` + +3. To further scale up your deployment with autoscaling and failure recovery, check out the [SkyServe + SGLang guide](https://github.com/skypilot-org/skypilot/tree/master/llm/sglang#serving-llama-2-with-sglang-for-more-traffic-using-skyserve). +
+ +## Common Notes + +- [FlashInfer](https://github.com/flashinfer-ai/flashinfer) is the default attention kernel backend. It only supports sm75 and above. If you encounter any FlashInfer-related issues on sm75+ devices (e.g., T4, A10, A100, L4, L40S, H100), please switch to other kernels by adding `--attention-backend triton --sampling-backend pytorch` and open an issue on GitHub. +- To reinstall flashinfer locally, use the following command: `pip3 install --upgrade flashinfer-python --force-reinstall --no-deps` and then delete the cache with `rm -rf ~/.cache/flashinfer`. +- If you only need to use OpenAI API models with the frontend language, you can avoid installing other dependencies by using `pip install "sglang[openai]"`. +- The language frontend operates independently of the backend runtime. You can install the frontend locally without needing a GPU, while the backend can be set up on a GPU-enabled machine. To install the frontend, run `pip install sglang`, and for the backend, use `pip install sglang[srt]`. `srt` is the abbreviation of SGLang runtime. diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 000000000..f948fca24 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,95 @@ +SGLang Documentation +==================== + +SGLang is a fast serving framework for large language models and vision language models. +It makes your interaction with models faster and more controllable by co-designing the backend runtime and frontend language. +The core features include: + +- **Fast Backend Runtime**: Provides efficient serving with RadixAttention for prefix caching, zero-overhead CPU scheduler, prefill-decode disaggregation, speculative decoding, continuous batching, paged attention, tensor/pipeline/expert/data parallelism, structured outputs, chunked prefill, quantization (FP4/FP8/INT4/AWQ/GPTQ), and multi-lora batching. +- **Flexible Frontend Language**: Offers an intuitive interface for programming LLM applications, including chained generation calls, advanced prompting, control flow, multi-modal inputs, parallelism, and external interactions. +- **Extensive Model Support**: Supports a wide range of generative models (Llama, Qwen, DeepSeek, Kimi, GPT, Gemma, Mistral, etc.), embedding models (e5-mistral, gte, mcdse) and reward models (Skywork), with easy extensibility for integrating new models. +- **Active Community**: SGLang is open-source and backed by an active community with wide industry adoption. + +.. toctree:: + :maxdepth: 1 + :caption: Get Started + + get_started/install.md + +.. toctree:: + :maxdepth: 1 + :caption: Basic Usage + + basic_usage/send_request.ipynb + basic_usage/openai_api.rst + basic_usage/offline_engine_api.ipynb + basic_usage/native_api.ipynb + basic_usage/sampling_params.md + basic_usage/deepseek.md + basic_usage/gpt_oss.md + basic_usage/llama4.md + basic_usage/qwen3.md + +.. toctree:: + :maxdepth: 1 + :caption: Advanced Features + + advanced_features/server_arguments.md + advanced_features/hyperparameter_tuning.md + advanced_features/speculative_decoding.ipynb + advanced_features/structured_outputs.ipynb + advanced_features/structured_outputs_for_reasoning_models.ipynb + advanced_features/tool_parser.ipynb + advanced_features/separate_reasoning.ipynb + advanced_features/quantization.md + advanced_features/lora.ipynb + advanced_features/pd_disaggregation.md + advanced_features/vlm_query.ipynb + advanced_features/router.md + advanced_features/observability.md + advanced_features/attention_backend.md + +.. toctree:: + :maxdepth: 1 + :caption: Supported Models + + supported_models/generative_models.md + supported_models/multimodal_language_models.md + supported_models/embedding_models.md + supported_models/reward_models.md + supported_models/rerank_models.md + supported_models/support_new_models.md + supported_models/transformers_fallback.md + supported_models/modelscope.md + +.. toctree:: + :maxdepth: 1 + :caption: Hardware Platforms + + platforms/amd_gpu.md + platforms/blackwell_gpu.md + platforms/cpu_server.md + platforms/tpu.md + platforms/nvidia_jetson.md + platforms/ascend_npu.md + +.. toctree:: + :maxdepth: 1 + :caption: Developer Guide + + developer_guide/contribution_guide.md + developer_guide/development_guide_using_docker.md + developer_guide/benchmark_and_profiling.md + developer_guide/bench_serving.md + +.. toctree:: + :maxdepth: 1 + :caption: References + + references/faq.md + references/environment_variables.md + references/production_metrics.md + references/multi_node_deployment/multi_node_index.rst + references/custom_chat_template.md + references/frontend/frontend_index.rst + references/learn_more.md diff --git a/docs/platforms/amd_gpu.md b/docs/platforms/amd_gpu.md new file mode 100644 index 000000000..81d6d544a --- /dev/null +++ b/docs/platforms/amd_gpu.md @@ -0,0 +1,158 @@ +# AMD GPUs + +This document describes how run SGLang on AMD GPUs. If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues). + +## System Configuration + +When using AMD GPUs (such as MI300X), certain system-level optimizations help ensure stable performance. Here we take MI300X as an example. AMD provides official documentation for MI300X optimization and system tuning: + +- [AMD MI300X Tuning Guides](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/index.html) +- [LLM inference performance validation on AMD Instinct MI300X](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference/vllm-benchmark.html) +- [AMD Instinct MI300X System Optimization](https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html) +- [AMD Instinct MI300X Workload Optimization](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference-optimization/workload.html) +- [Supercharge DeepSeek-R1 Inference on AMD Instinct MI300X](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1-Part2/README.html) + +**NOTE:** We strongly recommend reading these docs and guides entirely to fully utilize your system. + +Below are a few key settings to confirm or enable for SGLang: + +### Update GRUB Settings + +In `/etc/default/grub`, append the following to `GRUB_CMDLINE_LINUX`: + +```text +pci=realloc=off iommu=pt +``` + +Afterward, run `sudo update-grub` (or your distro’s equivalent) and reboot. + +### Disable NUMA Auto-Balancing + +```bash +sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing' +``` + +You can automate or verify this change using [this helpful script](https://github.com/ROCm/triton/blob/rocm_env/scripts/amd/env_check.sh). + +Again, please go through the entire documentation to confirm your system is using the recommended configuration. + +## Install SGLang + +You can install SGLang using one of the methods below. + +### Install from Source + +```bash +# Use the last release branch +git clone -b v0.5.2 https://github.com/sgl-project/sglang.git +cd sglang + +# Compile sgl-kernel +pip install --upgrade pip +cd sgl-kernel +python setup_rocm.py install + +# Install sglang python package +cd .. +pip install -e "python[all_hip]" +``` + +### Install Using Docker (Recommended) + +The docker images are available on Docker Hub at [lmsysorg/sglang](https://hub.docker.com/r/lmsysorg/sglang/tags), built from [Dockerfile.rocm](https://github.com/sgl-project/sglang/tree/main/docker). + +The steps below show how to build and use an image. + +1. Build the docker image. + If you use pre-built images, you can skip this step and replace `sglang_image` with the pre-built image names in the steps below. + + ```bash + docker build -t sglang_image -f Dockerfile.rocm . + ``` + +2. Create a convenient alias. + + ```bash + alias drun='docker run -it --rm --network=host --privileged --device=/dev/kfd --device=/dev/dri \ + --ipc=host --shm-size 16G --group-add video --cap-add=SYS_PTRACE \ + --security-opt seccomp=unconfined \ + -v $HOME/dockerx:/dockerx \ + -v /data:/data' + ``` + + If you are using RDMA, please note that: + - `--network host` and `--privileged` are required by RDMA. If you don't need RDMA, you can remove them. + - You may need to set `NCCL_IB_GID_INDEX` if you are using RoCE, for example: `export NCCL_IB_GID_INDEX=3`. + +3. Launch the server. + + **NOTE:** Replace `` below with your [huggingface hub token](https://huggingface.co/docs/hub/en/security-tokens). + + ```bash + drun -p 30000:30000 \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=" \ + sglang_image \ + python3 -m sglang.launch_server \ + --model-path NousResearch/Meta-Llama-3.1-8B \ + --host 0.0.0.0 \ + --port 30000 + ``` + +4. To verify the utility, you can run a benchmark in another terminal or refer to [other docs](https://docs.sglang.ai/backend/openai_api_completions.html) to send requests to the engine. + + ```bash + drun sglang_image \ + python3 -m sglang.bench_serving \ + --backend sglang \ + --dataset-name random \ + --num-prompts 4000 \ + --random-input 128 \ + --random-output 128 + ``` + +With your AMD system properly configured and SGLang installed, you can now fully leverage AMD hardware to power SGLang’s machine learning capabilities. + +## Examples + +### Running DeepSeek-V3 + +The only difference when running DeepSeek-V3 is in how you start the server. Here's an example command: + +```bash +drun -p 30000:30000 \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --ipc=host \ + --env "HF_TOKEN=" \ + sglang_image \ + python3 -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-V3 \ # <- here + --tp 8 \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 30000 +``` + +[Running DeepSeek-R1 on a single NDv5 MI300X VM](https://techcommunity.microsoft.com/blog/azurehighperformancecomputingblog/running-deepseek-r1-on-a-single-ndv5-mi300x-vm/4372726) could also be a good reference. + +### Running Llama3.1 + +Running Llama3.1 is nearly identical to running DeepSeek-V3. The only difference is in the model specified when starting the server, shown by the following example command: + +```bash +drun -p 30000:30000 \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --ipc=host \ + --env "HF_TOKEN=" \ + sglang_image \ + python3 -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ # <- here + --tp 8 \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 30000 +``` + +### Warmup Step + +When the server displays `The server is fired up and ready to roll!`, it means the startup is successful. diff --git a/docs/platforms/ascend_npu.md b/docs/platforms/ascend_npu.md new file mode 100644 index 000000000..f57d3fe95 --- /dev/null +++ b/docs/platforms/ascend_npu.md @@ -0,0 +1,206 @@ +# Ascend NPUs + +You can install SGLang using any of the methods below. Please go through `System Settings` section to ensure the clusters are roaring at max performance. Feel free to leave an issue [here at sglang](https://github.com/sgl-project/sglang/issues) if you encounter any issues or have any problems. + +## System Settings + +### CPU performance power scheme + +The default power scheme on Ascend hardware is `ondemand` which could affect performance, changing it to `performance` is recommended. + +```shell +echo performance | sudo tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor + +# Make sure changes are applied successfully +cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor # shows performance +``` + +### Disable NUMA balancing + +```shell +sudo sysctl -w kernel.numa_balancing=0 + +# Check +cat /proc/sys/kernel/numa_balancing # shows 0 +``` + +### Prevent swapping out system memory + +```shell +sudo sysctl -w vm.swappiness=10 + +# Check +cat /proc/sys/vm/swappiness # shows 10 +``` + +## Installing SGLang + +### Method 1: Installing from source with prerequisites + +#### Python Version + +Only `python==3.11` is supported currently. If you don't want to break system pre-installed python, try installing with [conda](https://github.com/conda/conda). + +```shell +conda create --name sglang_npu python=3.11 +conda activate sglang_npu +``` + +#### MemFabric Adaptor + +_TODO: MemFabric is still a working project yet open sourced til August/September, 2025. We will release it as prebuilt wheel package for now._ + +_Notice: Prebuilt wheel package is based on `aarch64`, please leave an issue [here at sglang](https://github.com/sgl-project/sglang/issues) to let us know the requests for `amd64` build._ + +MemFabric Adaptor is a drop-in replacement of Mooncake Transfer Engine that enables KV cache transfer on Ascend NPU clusters. + +```shell +MF_WHL_NAME="mf_adapter-1.0.0-cp311-cp311-linux_aarch64.whl" +MEMFABRIC_URL="https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/${MF_WHL_NAME}" +wget -O "${MF_WHL_NAME}" "${MEMFABRIC_URL}" && pip install "./${MF_WHL_NAME}" +``` + +#### Pytorch and Pytorch Framework Adaptor on Ascend + +Only `torch==2.6.0` is supported currently due to NPUgraph and Triton-on-Ascend's limitation, however a more generalized version will be release by the end of September, 2025. + +```shell +PYTORCH_VERSION=2.6.0 +TORCHVISION_VERSION=0.21.0 +pip install torch==$PYTORCH_VERSION torchvision==$TORCHVISION_VERSION --index-url https://download.pytorch.org/whl/cpu + +PTA_VERSION="v7.1.0.1-pytorch2.6.0" +PTA_NAME="torch_npu-2.6.0.post1-cp311-cp311-manylinux_2_28_aarch64.whl" +PTA_URL="https://gitee.com/ascend/pytorch/releases/download/${PTA_VERSION}/${PTA_WHL_NAME}" +wget -O "${PTA_NAME}" "${PTA_URL}" && pip install "./${PTA_NAME}" +``` + +#### vLLM + +vLLM is still a major prerequisite on Ascend NPU. Because of `torch==2.6.0` limitation, only vLLM v0.8.5 is supported. + +```shell +VLLM_TAG=v0.8.5 +git clone --depth 1 https://github.com/vllm-project/vllm.git --branch $VLLM_TAG +(cd vllm && VLLM_TARGET_DEVICE="empty" pip install -v -e .) +``` + +#### Triton on Ascend + +_Notice:_ We recommend installing triton-ascend from source due to its rapid development, the version on PYPI can't keep up for know. This problem will be solved on Sep. 2025, afterwards `pip install` would be the one and only installing method. + +Please follow Triton-on-Ascend's [installation guide from source](https://gitee.com/ascend/triton-ascend#2%E6%BA%90%E4%BB%A3%E7%A0%81%E5%AE%89%E8%A3%85-triton-ascend) to install the latest `triton-ascend` package. + +#### DeepEP-compatible Library + +We are also providing a DeepEP-compatible Library as a drop-in replacement of deepseek-ai's DeepEP library, check the [installation guide](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md). + +#### Installing SGLang from source + +```shell +# Use the last release branch +git clone -b v0.5.2 https://github.com/sgl-project/sglang.git +cd sglang + +pip install --upgrade pip +pip install -e python[srt_npu] +``` + +### Method 2: Using docker + +__Notice:__ `--privileged` and `--network=host` are required by RDMA, which is typically needed by Ascend NPU clusters. + +__Notice:__ The following docker command is based on Atlas 800I A3 machines. If you are using Atlas 800I A2, make sure only `davinci[0-7]` are mapped into container. + +```shell +# Clone the SGLang repository +git clone https://github.com/sgl-project/sglang.git +cd sglang/docker + +# Build the docker image +docker build -t sglang-npu:main -f Dockerfile.npu . + +alias drun='docker run -it --rm --privileged --network=host --ipc=host --shm-size=16g \ + --device=/dev/davinci0 --device=/dev/davinci1 --device=/dev/davinci2 --device=/dev/davinci3 \ + --device=/dev/davinci4 --device=/dev/davinci5 --device=/dev/davinci6 --device=/dev/davinci7 \ + --device=/dev/davinci8 --device=/dev/davinci9 --device=/dev/davinci10 --device=/dev/davinci11 \ + --device=/dev/davinci12 --device=/dev/davinci13 --device=/dev/davinci14 --device=/dev/davinci15 \ + --device=/dev/davinci_manager --device=/dev/hisi_hdc \ + --volume /usr/local/sbin:/usr/local/sbin --volume /usr/local/Ascend/driver:/usr/local/Ascend/driver \ + --volume /usr/local/Ascend/firmware:/usr/local/Ascend/firmware \ + --volume /etc/ascend_install.info:/etc/ascend_install.info \ + --volume /var/queue_schedule:/var/queue_schedule --volume ~/.cache/:/root/.cache/' + +drun --env "HF_TOKEN=" \ + sglang-npu:main \ + python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --attention-backend ascend --host 0.0.0.0 --port 30000 +``` + +## Examples + +### Running DeepSeek-V3 + +Running DeepSeek with PD disaggregation on 2 x Atlas 800I A3. +Model weights could be found [here](https://modelers.cn/models/State_Cloud/Deepseek-R1-bf16-hfd-w8a8). + +Prefill: + +```shell +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export ASCEND_MF_STORE_URL="tcp://:" + +drun sglang-npu:main \ + python3 -m sglang.launch_server --model-path State_Cloud/DeepSeek-R1-bf16-hfd-w8a8 \ + --trust-remote-code \ + --attention-backend ascend \ + --mem-fraction-static 0.8 \ + --quantization w8a8_int8 \ + --tp-size 16 \ + --dp-size 1 \ + --nnodes 1 \ + --node-rank 0 \ + --disaggregation-mode prefill \ + --disaggregation-bootstrap-port 6657 \ + --disaggregation-transfer-backend ascend \ + --dist-init-addr :6688 \ + --host \ + --port 8000 +``` + +Decode: + +```shell +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export ASCEND_MF_STORE_URL="tcp://:" +export HCCL_BUFFSIZE=200 +export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=24 + +drun sglang-npu:main \ + python3 -m sglang.launch_server --model-path State_Cloud/DeepSeek-R1-bf16-hfd-w8a8 \ + --trust-remote-code \ + --attention-backend ascend \ + --mem-fraction-static 0.8 \ + --quantization w8a8_int8 \ + --enable-deepep-moe \ + --deepep-mode low_latency \ + --tp-size 16 \ + --dp-size 1 \ + --ep-size 16 \ + --nnodes 1 \ + --node-rank 0 \ + --disaggregation-mode decode \ + --disaggregation-transfer-backend ascend \ + --dist-init-addr :6688 \ + --host \ + --port 8001 +``` + +Mini_LB: + +```shell +drun sglang-npu:main \ + python -m sglang.srt.disaggregation.launch_lb \ + --prefill http://:8000 \ + --decode http://:8001 \ + --host 127.0.0.1 --port 5000 +``` diff --git a/docs/platforms/blackwell_gpu.md b/docs/platforms/blackwell_gpu.md new file mode 100644 index 000000000..8c433b3f0 --- /dev/null +++ b/docs/platforms/blackwell_gpu.md @@ -0,0 +1,9 @@ +# Blackwell GPUs + +We will release the pre-built wheels soon. Before that, please try to compile from source or check the blackwell docker images from [lmsysorg/sglang](https://hub.docker.com/r/lmsysorg/sglang/tags). + +## B200 with x86 CPUs +TODO + +## GB200/GB300 with ARM CPUs +TODO diff --git a/docs/platforms/cpu_server.md b/docs/platforms/cpu_server.md new file mode 100644 index 000000000..97fad918d --- /dev/null +++ b/docs/platforms/cpu_server.md @@ -0,0 +1,204 @@ +# CPU Servers + +The document addresses how to set up the [SGLang](https://github.com/sgl-project/sglang) environment and run LLM inference on CPU servers. +Specifically, SGLang is well optimized on the CPUs equipped with Intel® AMX® Instructions, +which are 4th generation or newer Intel® Xeon® Scalable Processors. + +## Optimized Model List + +A list of popular LLMs are optimized and run efficiently on CPU, +including the most notable open-source models like Llama series, Qwen series, +and the phenomenal high-quality reasoning model DeepSeek-R1. + +| Model Name | BF16 | w8a8_int8 | FP8 | +|:---:|:---:|:---:|:---:| +| DeepSeek-R1 | | [meituan/DeepSeek-R1-Channel-INT8](https://huggingface.co/meituan/DeepSeek-R1-Channel-INT8) | [deepseek-ai/DeepSeek-R1](https://huggingface.co/deepseek-ai/DeepSeek-R1) | +| Llama-3.2-3B | [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) | [RedHatAI/Llama-3.2-3B-quantized.w8a8](https://huggingface.co/RedHatAI/Llama-3.2-3B-Instruct-quantized.w8a8) | | +| Llama-3.1-8B | [meta-llama/Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) | [RedHatAI/Meta-Llama-3.1-8B-quantized.w8a8](https://huggingface.co/RedHatAI/Meta-Llama-3.1-8B-quantized.w8a8) | | +| QwQ-32B | | [RedHatAI/QwQ-32B-quantized.w8a8](https://huggingface.co/RedHatAI/QwQ-32B-quantized.w8a8) | | +| DeepSeek-Distilled-Llama | | [RedHatAI/DeepSeek-R1-Distill-Llama-70B-quantized.w8a8](https://huggingface.co/RedHatAI/DeepSeek-R1-Distill-Llama-70B-quantized.w8a8) | | +| Qwen3-235B | | | [Qwen/Qwen3-235B-A22B-FP8](https://huggingface.co/Qwen/Qwen3-235B-A22B-FP8) | + +**Note:** The model identifiers listed in the table above +have been verified on 6th Gen Intel® Xeon® P-core platforms. + +## Installation + +### Install Using Docker + +It is recommended to use Docker for setting up the SGLang environment. +A [Dockerfile](https://github.com/sgl-project/sglang/blob/main/docker/Dockerfile.xeon) is provided to facilitate the installation. +Replace `` below with your [HuggingFace access token](https://huggingface.co/docs/hub/en/security-tokens). + +```bash +# Clone the SGLang repository +git clone https://github.com/sgl-project/sglang.git +cd sglang/docker + +# Build the docker image +docker build -t sglang-cpu:main -f Dockerfile.xeon . + +# Initiate a docker container +docker run \ + -it \ + --privileged \ + --ipc=host \ + --network=host \ + -v /dev/shm:/dev/shm \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + -p 30000:30000 \ + -e "HF_TOKEN=" \ + sglang-cpu:main /bin/bash +``` + +### Install From Source + +If you'd prefer to install SGLang in a bare metal environment, +the command list is as below. +It is worth noting that the environment variable `SGLANG_USE_CPU_ENGINE=1` +is required to enable SGLang service with CPU engine. + +```bash +# Create and activate a conda environment +conda create -n sgl-cpu python=3.12 -y +conda activate sgl-cpu + +# Optional: Set PyTorch CPU as primary pip install channel to avoid installing CUDA version +pip config set global.index-url https://download.pytorch.org/whl/cpu +pip config set global.extra-index-url https://pypi.org/simple + +# Check if some conda related environment variables have been set +env | grep -i conda +# The following environment variable settings are required +# if they have not been set properly +export CONDA_EXE=$(which conda) +export CONDA_ROOT=${CONDA_EXE}/../.. +export CONDA_PREFIX=${CONDA_ROOT}/envs/sgl-cpu +export PATH=${PATH}:${CONDA_ROOT}/bin:${CONDA_ROOT}/condabin + +# Clone the SGLang code +git clone https://github.com/sgl-project/sglang.git +cd sglang +git checkout + +# Install SGLang dependent libs, and build SGLang main package +pip install --upgrade pip setuptools +conda install -y libsqlite==3.48.0 gperftools tbb libnuma numactl +pip install -e "python[all_cpu]" +pip install torch==2.7.1 torchvision==0.22.1 triton==3.3.1 --force-reinstall + +# Build the CPU backend kernels +cd sgl-kernel +cp pyproject_cpu.toml pyproject.toml +pip install . + +# Other required environment variables +# Recommend to set these in ~/.bashrc in order not to set every time in a new terminal +export SGLANG_USE_CPU_ENGINE=1 +export LD_PRELOAD=${LD_PRELOAD}:${CONDA_PREFIX}/lib/libiomp5.so:${CONDA_PREFIX}/lib/libtcmalloc.so:${CONDA_PREFIX}/lib/libtbbmalloc.so.2 +``` + +## Launch of the Serving Engine + +Example command to launch SGLang serving: + +```bash +python -m sglang.launch_server \ + --model \ + --trust-remote-code \ + --disable-overlap-schedule \ + --device cpu \ + --host 0.0.0.0 \ + --tp 6 +``` + +Notes: + +1. For running W8A8 quantized models, please add the flag `--quantization w8a8_int8`. + +2. The flag `--tp 6` specifies that tensor parallelism will be applied using 6 ranks (TP6). + The number of TP specified is how many TP ranks will be used during the execution. + In a CPU platform, a TP rank means a sub-NUMA cluster (SNC). + Usually we can get the SNC information (How many available) from Operation System. + User can specify TP to be no more than the total available SNCs in current system. + + If the specified TP rank number differs from the total SNC count, + the system will automatically utilize the first `n` SNCs. + Note that `n` cannot exceed the total SNC number, doing so will result in an error. + + To specify the cores to be used, we need to explicitly set the environment variable `SGLANG_CPU_OMP_THREADS_BIND`. + For example, if we want to run the SGLang service using the first 40 cores of each SNC on a Xeon® 6980P server, + which has 43-43-42 cores on the 3 SNCs of a socket, we should set: + + ```bash + export SGLANG_CPU_OMP_THREADS_BIND="0-39|43-82|86-125|128-167|171-210|214-253" + ``` + + Please beware that with SGLANG_CPU_OMP_THREADS_BIND set, + the available memory amounts of the ranks may not be determined in prior. + You may need to set proper `--max-total-tokens` to avoid the out-of-memory error. + +3. For optimizing decoding with torch.compile, please add the flag `--enable-torch-compile`. + To specify the maximum batch size when using torch compile, set the flag `--torch-compile-max-bs`. + For example, `--enable-torch-compile --torch-compile-max-bs 4` means using torch compile and setting the + maximum batch size to 4. + +4. A warmup step is automatically triggered when the service is started. + The server is ready when you see the log `The server is fired up and ready to roll!`. + +## Benchmarking with Requests + +You can benchmark the performance via the `bench_serving` script. +Run the command in another terminal. + +```bash +python -m sglang.bench_serving \ + --dataset-name random \ + --random-input-len 1024 \ + --random-output-len 1024 \ + --num-prompts 1 \ + --request-rate inf \ + --random-range-ratio 1.0 +``` + +The detail explanations of the parameters can be looked up by the command: + +```bash +python -m sglang.bench_serving -h +``` + +Additionally, the requests can be formed with +[OpenAI Completions API](https://docs.sglang.ai/basic_usage/openai_api_completions.html) +and sent via the command line (e.g. using `curl`) or via your own script. + +## Example: Running DeepSeek-R1 + +An example command to launch service for W8A8 DeepSeek-R1 on a Xeon® 6980P server + +```bash +python -m sglang.launch_server \ + --model meituan/DeepSeek-R1-Channel-INT8 \ + --trust-remote-code \ + --disable-overlap-schedule \ + --device cpu \ + --quantization w8a8_int8 \ + --host 0.0.0.0 \ + --mem-fraction-static 0.8 \ + --tp 6 +``` + +Similarly, an example command to launch service for FP8 DeepSeek-R1 would be + +```bash +python -m sglang.launch_server \ + --model deepseek-ai/DeepSeek-R1 \ + --trust-remote-code \ + --disable-overlap-schedule \ + --device cpu \ + --host 0.0.0.0 \ + --mem-fraction-static 0.8 \ + --tp 6 +``` + +Then you can test with `bench_serving` command or construct your own command or script +following [the benchmarking example](#benchmarking-with-requests). diff --git a/docs/platforms/nvidia_jetson.md b/docs/platforms/nvidia_jetson.md new file mode 100644 index 000000000..362f60c83 --- /dev/null +++ b/docs/platforms/nvidia_jetson.md @@ -0,0 +1,80 @@ +# NVIDIA Jetson Orin + +## Prerequisites + +Before starting, ensure the following: + +- [**NVIDIA Jetson AGX Orin Devkit**](https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-orin/) is set up with **JetPack 6.1** or later. +- **CUDA Toolkit** and **cuDNN** are installed. +- Verify that the Jetson AGX Orin is in **high-performance mode**: +```bash +sudo nvpmodel -m 0 +``` +* * * * * +## Installing and running SGLang with Jetson Containers +Clone the jetson-containers github repository: +``` +git clone https://github.com/dusty-nv/jetson-containers.git +``` +Run the installation script: +``` +bash jetson-containers/install.sh +``` +Build the container image: +``` +jetson-containers build sglang +``` +Run the container: +``` +jetson-containers run $(autotag sglang) +``` +Or you can also manually run a container with this command: +``` +docker run --runtime nvidia -it --rm --network=host IMAGE_NAME +``` +* * * * * + +Running Inference +----------------------------------------- + +Launch the server: +```bash +python -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-R1-Distill-Llama-8B \ + --device cuda \ + --dtype half \ + --attention-backend flashinfer \ + --mem-fraction-static 0.8 \ + --context-length 8192 +``` +The quantization and limited context length (`--dtype half --context-length 8192`) are due to the limited computational resources in [Nvidia jetson kit](https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-orin/). A detailed explanation can be found in [Server Arguments](../backend/server_arguments.md). + +After launching the engine, refer to [Chat completions](https://docs.sglang.ai/backend/openai_api_completions.html#Usage) to test the usability. +* * * * * +Running quantization with TorchAO +------------------------------------- +TorchAO is suggested to NVIDIA Jetson Orin. +```bash +python -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --device cuda \ + --dtype bfloat16 \ + --attention-backend flashinfer \ + --mem-fraction-static 0.8 \ + --context-length 8192 \ + --torchao-config int4wo-128 +``` +This enables TorchAO's int4 weight-only quantization with a 128-group size. The usage of `--torchao-config int4wo-128` is also for memory efficiency. + + +* * * * * +Structured output with XGrammar +------------------------------- +Please refer to [SGLang doc structured output](../advanced_features/structured_outputs.ipynb). +* * * * * + +Thanks to the support from [Nurgaliyev Shakhizat](https://github.com/shahizat), [Dustin Franklin](https://github.com/dusty-nv) and [Johnny Núñez Cano](https://github.com/johnnynunez). + +References +---------- +- [NVIDIA Jetson AGX Orin Documentation](https://developer.nvidia.com/embedded/jetson-agx-orin) diff --git a/docs/platforms/tpu.md b/docs/platforms/tpu.md new file mode 100644 index 000000000..f304234cf --- /dev/null +++ b/docs/platforms/tpu.md @@ -0,0 +1,3 @@ +# TPU + +The support for TPU is under active development. Please stay tuned. diff --git a/docs/references/custom_chat_template.md b/docs/references/custom_chat_template.md new file mode 100644 index 000000000..557af5bf5 --- /dev/null +++ b/docs/references/custom_chat_template.md @@ -0,0 +1,42 @@ +# Custom Chat Template + +**NOTE**: There are two chat template systems in SGLang project. This document is about setting a custom chat template for the OpenAI-compatible API server (defined at [conversation.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/conversation.py)). It is NOT related to the chat template used in the SGLang language frontend (defined at [chat_template.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/lang/chat_template.py)). + +By default, the server uses the chat template specified in the model tokenizer from Hugging Face. +It should just work for most official models such as Llama-2/Llama-3. + +If needed, you can also override the chat template when launching the server: + +```bash +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --chat-template llama-2 +``` + +If the chat template you are looking for is missing, you are welcome to contribute it or load it from a file. + +## JSON Format + +You can load the JSON format, which is defined by `conversation.py`. + +```json +{ + "name": "my_model", + "system": "<|im_start|>system", + "user": "<|im_start|>user", + "assistant": "<|im_start|>assistant", + "sep_style": "CHATML", + "sep": "<|im_end|>", + "stop_str": ["<|im_end|>", "<|im_start|>"] +} +``` + +```bash +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --chat-template ./my_model_template.json +``` + +## Jinja Format + +You can also use the [Jinja template format](https://huggingface.co/docs/transformers/main/en/chat_templating) as defined by Hugging Face Transformers. + +```bash +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --chat-template ./my_model_template.jinja +``` diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md new file mode 100644 index 000000000..f22685454 --- /dev/null +++ b/docs/references/environment_variables.md @@ -0,0 +1,97 @@ +# Environment Variables + +SGLang supports various environment variables that can be used to configure its runtime behavior. This document provides a comprehensive list and aims to stay updated over time. + +*Note: SGLang uses two prefixes for environment variables: `SGL_` and `SGLANG_`. This is likely due to historical reasons. While both are currently supported for different settings, future versions might consolidate them.* + +## General Configuration + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_USE_MODELSCOPE` | Enable using models from ModelScope | `false` | +| `SGLANG_HOST_IP` | Host IP address for the server | `0.0.0.0` | +| `SGLANG_PORT` | Port for the server | auto-detected | +| `SGLANG_LOGGING_CONFIG_PATH` | Custom logging configuration path | Not set | +| `SGLANG_DISABLE_REQUEST_LOGGING` | Disable request logging | `false` | +| `SGLANG_HEALTH_CHECK_TIMEOUT` | Timeout for health check in seconds | `20` | + +## Performance Tuning + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_ENABLE_TORCH_INFERENCE_MODE` | Control whether to use torch.inference_mode | `false` | +| `SGLANG_ENABLE_TORCH_COMPILE` | Enable torch.compile | `true` | +| `SGLANG_SET_CPU_AFFINITY` | Enable CPU affinity setting (often set to `1` in Docker builds) | `0` | +| `SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN` | Allows the scheduler to overwrite longer context length requests (often set to `1` in Docker builds) | `0` | +| `SGLANG_IS_FLASHINFER_AVAILABLE` | Control FlashInfer availability check | `true` | +| `SGLANG_SKIP_P2P_CHECK` | Skip P2P (peer-to-peer) access check | `false` | +| `SGL_CHUNKED_PREFIX_CACHE_THRESHOLD` | Sets the threshold for enabling chunked prefix caching | `8192` | +| `SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION` | Enable RoPE fusion in Fused Multi-Layer Attention | `1` | + +## DeepGEMM Configuration (Advanced Optimization) + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGL_ENABLE_JIT_DEEPGEMM` | Enable Just-In-Time compilation of DeepGEMM kernels | `"true"` | +| `SGL_JIT_DEEPGEMM_PRECOMPILE` | Enable precompilation of DeepGEMM kernels | `"true"` | +| `SGL_JIT_DEEPGEMM_COMPILE_WORKERS` | Number of workers for parallel DeepGEMM kernel compilation | `4` | +| `SGL_IN_DEEPGEMM_PRECOMPILE_STAGE` | Indicator flag used during the DeepGEMM precompile script | `"false"` | +| `SGL_DG_CACHE_DIR` | Directory for caching compiled DeepGEMM kernels | `~/.cache/deep_gemm` | +| `SGL_DG_USE_NVRTC` | Use NVRTC (instead of Triton) for JIT compilation (Experimental) | `"0"` | +| `SGL_USE_DEEPGEMM_BMM` | Use DeepGEMM for Batched Matrix Multiplication (BMM) operations | `"false"` | + +## Memory Management + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_DEBUG_MEMORY_POOL` | Enable memory pool debugging | `false` | +| `SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION` | Clip max new tokens estimation for memory planning | `4096` | +| `SGLANG_DETOKENIZER_MAX_STATES` | Maximum states for detokenizer | Default value based on system | +| `SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK` | Disable checks for memory imbalance across Tensor Parallel ranks | Not set (defaults to enabled check) | + +## Model-Specific Options + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_USE_AITER` | Use AITER optimize implementation | `false` | +| `SGLANG_INT4_WEIGHT` | Enable INT4 weight quantization | `false` | +| `SGLANG_MOE_PADDING` | Enable MoE padding (sets padding size to 128 if value is `1`, often set to `1` in Docker builds) | `0` | +| `SGLANG_FORCE_FP8_MARLIN` | Force using FP8 MARLIN kernels even if other FP8 kernels are available | `false` | +| `SGLANG_ENABLE_FLASHINFER_GEMM` | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs | `false` | +| `SGLANG_SUPPORT_CUTLASS_BLOCK_FP8` | Use Cutlass kernels when running blockwise fp8 GEMM on Hopper or Blackwell GPUs | `false` | +| `SGLANG_CUTLASS_MOE` | Use Cutlass FP8 MoE kernel on Blackwell GPUs | `false` | + + +## Distributed Computing + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_BLOCK_NONZERO_RANK_CHILDREN` | Control blocking of non-zero rank children processes | `1` | +| `SGL_IS_FIRST_RANK_ON_NODE` | Indicates if the current process is the first rank on its node | `"true"` | +| `SGLANG_PP_LAYER_PARTITION` | Pipeline parallel layer partition specification | Not set | + +## Testing & Debugging (Internal/CI) + +*These variables are primarily used for internal testing, continuous integration, or debugging.* + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_IS_IN_CI` | Indicates if running in CI environment | `false` | +| `SGLANG_AMD_CI` | Indicates running in AMD CI environment | `0` | +| `SGLANG_TEST_RETRACT` | Enable retract decode testing | `false` | +| `SGLANG_RECORD_STEP_TIME` | Record step time for profiling | `false` | +| `SGLANG_TEST_REQUEST_TIME_STATS` | Test request time statistics | `false` | +| `SGLANG_CI_SMALL_KV_SIZE` | Use small KV cache size in CI | Not set | + +## Profiling & Benchmarking + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_TORCH_PROFILER_DIR` | Directory for PyTorch profiler output | `/tmp` | +| `SGLANG_PROFILE_WITH_STACK` | Set `with_stack` option (bool) for PyTorch profiler (capture stack trace) | `true` | + +## Storage & Caching + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_DISABLE_OUTLINES_DISK_CACHE` | Disable Outlines disk cache | `true` | diff --git a/docs/references/faq.md b/docs/references/faq.md new file mode 100644 index 000000000..6d575d253 --- /dev/null +++ b/docs/references/faq.md @@ -0,0 +1,35 @@ +# Troubleshooting and Frequently Asked Questions + +## Troubleshooting + +This page lists common errors and tips for resolving them. + +### CUDA Out of Memory +If you encounter out-of-memory (OOM) errors, you can adjust the following parameters: + +- If OOM occurs during prefill, try reducing `--chunked-prefill-size` to `4096` or `2048`. This saves memory but slows down the prefill speed for long prompts. +- If OOM occurs during decoding, try lowering `--max-running-requests`. +- You can also reduce `--mem-fraction-static` to a smaller value, such as 0.8 or 0.7. This decreases the memory usage of the KV cache memory pool and helps prevent OOM errors during both prefill and decoding. However, it limits maximum concurrency and reduces peak throughput. +- Another common case for OOM is requesting input logprobs for a long prompt as it requires significant memory. To address this, set `logprob_start_len` in your sampling parameters to include only the necessary parts. If you do need input logprobs for a long prompt, try reducing `--mem-fraction-static`. + +### CUDA Error: Illegal Memory Access Encountered +This error may result from kernel errors or out-of-memory issues: +- If it is a kernel error, resolving it may be challenging. Please file an issue on GitHub. +- If it is an out-of-memory issue, it may sometimes be reported as this error instead of "Out of Memory." Refer to the section above for guidance on avoiding OOM issues. + + +## Frequently Asked Questions + +### The results are not deterministic, even with a temperature of 0 + +You may notice that when you send the same request twice, the results from the engine will be slightly different, even when the temperature is set to 0. + +From our initial investigation, this indeterminism arises from two factors: dynamic batching and prefix caching. Roughly speaking, dynamic batching accounts for about 95% of the indeterminism, while prefix caching accounts for the remaining portion. The server runs dynamic batching under the hood. Different batch sizes can cause PyTorch/CuBLAS to dispatch to different CUDA kernels, which can lead to slight numerical differences. This difference accumulates across many layers, resulting in nondeterministic output when the batch size changes. Similarly, when prefix caching is enabled, it can also dispatch to different kernels. Even when the computations are mathematically equivalent, small numerical differences from different kernel implementations lead to the final nondeterministic outputs. + +To achieve more deterministic outputs in the current code, you can add `--disable-radix-cache` and send only one request at a time. The results will be mostly deterministic under this setting. + +We are still investigating the root causes and potential solutions. In the short term, we may introduce a "deterministic mode" that uses more padding to address the variance caused by dynamic batching. This mode will be more deterministic but slower. + +We have two issues to track our progress: +- The deterministic mode is tracked at [https://github.com/sgl-project/sglang/issues/1729](https://github.com/sgl-project/sglang/issues/1729). +- The per-request random seed is tracked at [https://github.com/sgl-project/sglang/issues/1335](https://github.com/sgl-project/sglang/issues/1335). diff --git a/docs/references/frontend/choices_methods.md b/docs/references/frontend/choices_methods.md new file mode 100644 index 000000000..30a0a1814 --- /dev/null +++ b/docs/references/frontend/choices_methods.md @@ -0,0 +1,77 @@ +# Choices Methods in SGLang +This doc describes the choices methods supported by SGLang. + +The optional `choices_method` arg determines how options supplied to SGLang's `choices` primitive are selected. Only the `RuntimeEndpoint` backend supports the `choices_method` arg. Other backends, such as `OpenAI`, have bespoke selection implementations due to API limitations. + +## Methods + +### Token Length Normalized + +Token length normalized is the default SGLang choices method. It selects the option with the highest average logprob across all of its tokens. + +Usage example (alternatively, simply omit the `choices_method` arg): +```python +@sgl.function +def example(s): + s += sgl.user("What is the capital of France?") + s += sgl.assistant( + sgl.gen( + "answer", + choices=["London", "Paris", "Berlin"], + choices_method=sgl.token_length_normalized, + ) + ) +``` + + +This can perform poorly if an option contains many tokens, where its later tokens are predicted with high confidence based on its earlier tokens. For instance, even strong models will fail the above example if the specified options are `["Paris", "Antidisestablishmentarianism"]`. + +### Greedy Token Selection + +Greedy token selection simply selects the option with the highest logprob for its initial token. For overlapping options where one option is a subset of a longer option, the logprobs of the shorter option are extended using its average logprob for comparison against the longer option. + +Usage example: +```python +@sgl.function +def example(s): + s += sgl.user("What is the capital of France?") + s += sgl.assistant( + sgl.gen( + "answer", + choices=["London", "Paris", "Berlin"], + choices_method=sgl.greedy_token_selection, + ) + ) +``` + +This can perform poorly if an option misleads the model down a bad path based on an attractive initial token. For instance, greedy selection will result in an incorrect response for this example: +```python +@sgl.function +def us_president_example(s): + s += sgl.user("Name a US president.") + s += sgl.assistant( + sgl.gen( + "answer", + choices=["Donald Duck", "Millard Fillmore"], + choices_method=sgl.greedy_token_selection, + ) + ) +``` + +### Unconditional Likelihood Normalized + +Unconditional likelihood normalized selects the option with the highest average token logprob once normalized by the unconditional token logprobs, as described in [this EleutherAI blogpost](https://blog.eleuther.ai/multiple-choice-normalization/). This method incurs an additional LLM call to obtain the unconditional likelihoods. + +Usage example: +```python +@sgl.function +def example(s): + s += sgl.user("What is the capital of France?") + s += sgl.assistant( + sgl.gen( + "answer", + choices=["London", "Paris", "Berlin"], + choices_method=sgl.unconditional_likelihood_normalized, + ) + ) +``` diff --git a/docs/references/frontend/frontend_index.rst b/docs/references/frontend/frontend_index.rst new file mode 100644 index 000000000..62544cba5 --- /dev/null +++ b/docs/references/frontend/frontend_index.rst @@ -0,0 +1,9 @@ +Frontend Language +================= + +.. toctree:: + :maxdepth: 1 + :caption: Frontend Language + + frontend_tutorial.ipynb + choices_methods.md diff --git a/docs/references/frontend/frontend_tutorial.ipynb b/docs/references/frontend/frontend_tutorial.ipynb new file mode 100644 index 000000000..836cab627 --- /dev/null +++ b/docs/references/frontend/frontend_tutorial.ipynb @@ -0,0 +1,456 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# SGLang Frontend Language" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SGLang frontend language can be used to define simple and easy prompts in a convenient, structured way." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server\n", + "\n", + "Launch the server in your terminal and wait for it to initialize." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang import assistant_begin, assistant_end\n", + "from sglang import assistant, function, gen, system, user\n", + "from sglang import image\n", + "from sglang import RuntimeEndpoint\n", + "from sglang.lang.api import set_default_backend\n", + "from sglang.srt.utils import load_image\n", + "from sglang.test.doc_patch import launch_server_cmd\n", + "from sglang.utils import print_highlight, terminate_process, wait_for_server\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct --host 0.0.0.0 --log-level warning\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "print(f\"Server started on http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Set the default backend. Note: Besides the local server, you may use also `OpenAI` or other API endpoints." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "set_default_backend(RuntimeEndpoint(f\"http://localhost:{port}\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Basic Usage\n", + "\n", + "The most simple way of using SGLang frontend language is a simple question answer dialog between a user and an assistant." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def basic_qa(s, question):\n", + " s += system(f\"You are a helpful assistant than can answer questions.\")\n", + " s += user(question)\n", + " s += assistant(gen(\"answer\", max_tokens=512))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "state = basic_qa(\"List 3 countries and their capitals.\")\n", + "print_highlight(state[\"answer\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi-turn Dialog\n", + "\n", + "SGLang frontend language can also be used to define multi-turn dialogs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def multi_turn_qa(s):\n", + " s += system(f\"You are a helpful assistant than can answer questions.\")\n", + " s += user(\"Please give me a list of 3 countries and their capitals.\")\n", + " s += assistant(gen(\"first_answer\", max_tokens=512))\n", + " s += user(\"Please give me another list of 3 countries and their capitals.\")\n", + " s += assistant(gen(\"second_answer\", max_tokens=512))\n", + " return s\n", + "\n", + "\n", + "state = multi_turn_qa()\n", + "print_highlight(state[\"first_answer\"])\n", + "print_highlight(state[\"second_answer\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Control flow\n", + "\n", + "You may use any Python code within the function to define more complex control flows." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def tool_use(s, question):\n", + " s += assistant(\n", + " \"To answer this question: \"\n", + " + question\n", + " + \". I need to use a \"\n", + " + gen(\"tool\", choices=[\"calculator\", \"search engine\"])\n", + " + \". \"\n", + " )\n", + "\n", + " if s[\"tool\"] == \"calculator\":\n", + " s += assistant(\"The math expression is: \" + gen(\"expression\"))\n", + " elif s[\"tool\"] == \"search engine\":\n", + " s += assistant(\"The key word to search is: \" + gen(\"word\"))\n", + "\n", + "\n", + "state = tool_use(\"What is 2 * 2?\")\n", + "print_highlight(state[\"tool\"])\n", + "print_highlight(state[\"expression\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Parallelism\n", + "\n", + "Use `fork` to launch parallel prompts. Because `sgl.gen` is non-blocking, the for loop below issues two generation calls in parallel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def tip_suggestion(s):\n", + " s += assistant(\n", + " \"Here are two tips for staying healthy: \"\n", + " \"1. Balanced Diet. 2. Regular Exercise.\\n\\n\"\n", + " )\n", + "\n", + " forks = s.fork(2)\n", + " for i, f in enumerate(forks):\n", + " f += assistant(\n", + " f\"Now, expand tip {i+1} into a paragraph:\\n\"\n", + " + gen(\"detailed_tip\", max_tokens=256, stop=\"\\n\\n\")\n", + " )\n", + "\n", + " s += assistant(\"Tip 1:\" + forks[0][\"detailed_tip\"] + \"\\n\")\n", + " s += assistant(\"Tip 2:\" + forks[1][\"detailed_tip\"] + \"\\n\")\n", + " s += assistant(\n", + " \"To summarize the above two tips, I can say:\\n\" + gen(\"summary\", max_tokens=512)\n", + " )\n", + "\n", + "\n", + "state = tip_suggestion()\n", + "print_highlight(state[\"summary\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Constrained Decoding\n", + "\n", + "Use `regex` to specify a regular expression as a decoding constraint. This is only supported for local models." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def regular_expression_gen(s):\n", + " s += user(\"What is the IP address of the Google DNS servers?\")\n", + " s += assistant(\n", + " gen(\n", + " \"answer\",\n", + " temperature=0,\n", + " regex=r\"((25[0-5]|2[0-4]\\d|[01]?\\d\\d?).){3}(25[0-5]|2[0-4]\\d|[01]?\\d\\d?)\",\n", + " )\n", + " )\n", + "\n", + "\n", + "state = regular_expression_gen()\n", + "print_highlight(state[\"answer\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use `regex` to define a `JSON` decoding schema." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "character_regex = (\n", + " r\"\"\"\\{\\n\"\"\"\n", + " + r\"\"\" \"name\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", + " + r\"\"\" \"house\": \"(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)\",\\n\"\"\"\n", + " + r\"\"\" \"blood status\": \"(Pure-blood|Half-blood|Muggle-born)\",\\n\"\"\"\n", + " + r\"\"\" \"occupation\": \"(student|teacher|auror|ministry of magic|death eater|order of the phoenix)\",\\n\"\"\"\n", + " + r\"\"\" \"wand\": \\{\\n\"\"\"\n", + " + r\"\"\" \"wood\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", + " + r\"\"\" \"core\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", + " + r\"\"\" \"length\": [0-9]{1,2}\\.[0-9]{0,2}\\n\"\"\"\n", + " + r\"\"\" \\},\\n\"\"\"\n", + " + r\"\"\" \"alive\": \"(Alive|Deceased)\",\\n\"\"\"\n", + " + r\"\"\" \"patronus\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", + " + r\"\"\" \"bogart\": \"[\\w\\d\\s]{1,16}\"\\n\"\"\"\n", + " + r\"\"\"\\}\"\"\"\n", + ")\n", + "\n", + "\n", + "@function\n", + "def character_gen(s, name):\n", + " s += user(\n", + " f\"{name} is a character in Harry Potter. Please fill in the following information about this character.\"\n", + " )\n", + " s += assistant(gen(\"json_output\", max_tokens=256, regex=character_regex))\n", + "\n", + "\n", + "state = character_gen(\"Harry Potter\")\n", + "print_highlight(state[\"json_output\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Batching \n", + "\n", + "Use `run_batch` to run a batch of prompts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def text_qa(s, question):\n", + " s += user(question)\n", + " s += assistant(gen(\"answer\", stop=\"\\n\"))\n", + "\n", + "\n", + "states = text_qa.run_batch(\n", + " [\n", + " {\"question\": \"What is the capital of the United Kingdom?\"},\n", + " {\"question\": \"What is the capital of France?\"},\n", + " {\"question\": \"What is the capital of Japan?\"},\n", + " ],\n", + " progress_bar=True,\n", + ")\n", + "\n", + "for i, state in enumerate(states):\n", + " print_highlight(f\"Answer {i+1}: {states[i]['answer']}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Streaming \n", + "\n", + "Use `stream` to stream the output to the user." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def text_qa(s, question):\n", + " s += user(question)\n", + " s += assistant(gen(\"answer\", stop=\"\\n\"))\n", + "\n", + "\n", + "state = text_qa.run(\n", + " question=\"What is the capital of France?\", temperature=0.1, stream=True\n", + ")\n", + "\n", + "for out in state.text_iter():\n", + " print(out, end=\"\", flush=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Complex Prompts\n", + "\n", + "You may use `{system|user|assistant}_{begin|end}` to define complex prompts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def chat_example(s):\n", + " s += system(\"You are a helpful assistant.\")\n", + " # Same as: s += s.system(\"You are a helpful assistant.\")\n", + "\n", + " with s.user():\n", + " s += \"Question: What is the capital of France?\"\n", + "\n", + " s += assistant_begin()\n", + " s += \"Answer: \" + gen(\"answer\", max_tokens=100, stop=\"\\n\")\n", + " s += assistant_end()\n", + "\n", + "\n", + "state = chat_example()\n", + "print_highlight(state[\"answer\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi-modal Generation\n", + "\n", + "You may use SGLang frontend language to define multi-modal prompts.\n", + "See [here](https://docs.sglang.ai/supported_models/generative_models.html) for supported models." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path Qwen/Qwen2.5-VL-7B-Instruct --host 0.0.0.0 --log-level warning\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "print(f\"Server started on http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "set_default_backend(RuntimeEndpoint(f\"http://localhost:{port}\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Ask a question about an image." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def image_qa(s, image_file, question):\n", + " s += user(image(image_file) + question)\n", + " s += assistant(gen(\"answer\", max_tokens=256))\n", + "\n", + "\n", + "image_url = \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + "image_bytes, _ = load_image(image_url)\n", + "state = image_qa(image_bytes, \"What is in the image?\")\n", + "print_highlight(state[\"answer\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/references/learn_more.md b/docs/references/learn_more.md new file mode 100644 index 000000000..b1a8a17da --- /dev/null +++ b/docs/references/learn_more.md @@ -0,0 +1,7 @@ +# Learn more + +You can find more blogs, slides, and videos about SGLang at [https://github.com/sgl-project/sgl-learning-materials](https://github.com/sgl-project/sgl-learning-materials). + +The latest SGLang features and updates are shared through the [LMSYS blog](https://lmsys.org/blog/). + +The 2025 H2 roadmap can be found at this [issue](https://github.com/sgl-project/sglang/issues/7736). diff --git a/docs/references/multi_node_deployment/deploy_on_k8s.md b/docs/references/multi_node_deployment/deploy_on_k8s.md new file mode 100644 index 000000000..cfc099f56 --- /dev/null +++ b/docs/references/multi_node_deployment/deploy_on_k8s.md @@ -0,0 +1,337 @@ +# Deploy On Kubernetes + +This document is for deploying a RoCE network-based SGLang two-node inference service on a Kubernetes (K8S) cluster. + +[LeaderWorkerSet (LWS)](https://github.com/kubernetes-sigs/lws) is a Kubernetes API that aims to address common deployment patterns of AI/ML inference workloads. A major use case is for multi-host/multi-node distributed inference. + +SGLang can also be deployed with LWS on Kubernetes for distributed model serving. + +Please see this guide for more details on deploying SGLang on Kubernetes using LWS. + +Here we take the deployment of DeepSeek-R1 as an example. + +## Prerequisites + +1. At least two Kubernetes nodes, each with two H20 systems and eight GPUs, are required. + +2. Make sure your K8S cluster has LWS correctly installed. If it hasn't been set up yet, please follow the [installation instructions](https://github.com/kubernetes-sigs/lws/blob/main/site/content/en/docs/installation/_index.md). **Note:** For LWS versions ≤0.5.x, you must use the Downward API to obtain `LWS_WORKER_INDEX`, as native support for this feature was introduced in v0.6.0. + +## Basic example + +For the basic example documentation, refer to [Deploy Distributed Inference Service with SGLang and LWS on GPUs](https://github.com/kubernetes-sigs/lws/tree/main/docs/examples/sglang). + +However, that document only covers the basic NCCL socket mode. + +In this section, we’ll make some simple modifications to adapt the setup to the RDMA scenario. + +## RDMA RoCE case + +* Check your env: + +```bash +[root@node1 ~]# ibstatus +Infiniband device 'mlx5_bond_0' port 1 status: + default gid: fe80:0000:0000:0000:0225:9dff:fe64:c79a + base lid: 0x0 + sm lid: 0x0 + state: 4: ACTIVE + phys state: 5: LinkUp + rate: 200 Gb/sec (2X NDR) + link_layer: Ethernet + +Infiniband device 'mlx5_bond_1' port 1 status: + default gid: fe80:0000:0000:0000:0225:9dff:fe6e:c3ec + base lid: 0x0 + sm lid: 0x0 + state: 4: ACTIVE + phys state: 5: LinkUp + rate: 200 Gb/sec (2X NDR) + link_layer: Ethernet + +Infiniband device 'mlx5_bond_2' port 1 status: + default gid: fe80:0000:0000:0000:0225:9dff:fe73:0dd7 + base lid: 0x0 + sm lid: 0x0 + state: 4: ACTIVE + phys state: 5: LinkUp + rate: 200 Gb/sec (2X NDR) + link_layer: Ethernet + +Infiniband device 'mlx5_bond_3' port 1 status: + default gid: fe80:0000:0000:0000:0225:9dff:fe36:f7ff + base lid: 0x0 + sm lid: 0x0 + state: 4: ACTIVE + phys state: 5: LinkUp + rate: 200 Gb/sec (2X NDR) + link_layer: Ethernet +``` + +* Prepare the `lws.yaml` file for deploying on k8s. + +```yaml +apiVersion: leaderworkerset.x-k8s.io/v1 +kind: LeaderWorkerSet +metadata: + name: sglang +spec: + replicas: 1 + leaderWorkerTemplate: + size: 2 + restartPolicy: RecreateGroupOnPodRestart + leaderTemplate: + metadata: + labels: + role: leader + spec: + dnsPolicy: ClusterFirstWithHostNet + hostNetwork: true + hostIPC: true + containers: + - name: sglang-leader + image: sglang:latest + securityContext: + privileged: true + env: + - name: NCCL_IB_GID_INDEX + value: "3" + command: + - python3 + - -m + - sglang.launch_server + - --model-path + - /work/models + - --mem-fraction-static + - "0.93" + - --torch-compile-max-bs + - "8" + - --max-running-requests + - "20" + - --tp + - "16" # Size of Tensor Parallelism + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20000 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --host + - "0.0.0.0" + - --port + - "40000" + resources: + limits: + nvidia.com/gpu: "8" + ports: + - containerPort: 40000 + readinessProbe: + tcpSocket: + port: 40000 + initialDelaySeconds: 15 + periodSeconds: 10 + volumeMounts: + - mountPath: /dev/shm + name: dshm + - name: model + mountPath: /work/models + - name: ib + mountPath: /dev/infiniband + volumes: + - name: dshm + emptyDir: + medium: Memory + - name: model + hostPath: + path: '< your models dir >' # modify it according your models dir + - name: ib + hostPath: + path: /dev/infiniband + workerTemplate: + spec: + dnsPolicy: ClusterFirstWithHostNet + hostNetwork: true + hostIPC: true + containers: + - name: sglang-worker + image: sglang:latest + securityContext: + privileged: true + env: + - name: NCCL_IB_GID_INDEX + value: "3" + command: + - python3 + - -m + - sglang.launch_server + - --model-path + - /work/models + - --mem-fraction-static + - "0.93" + - --torch-compile-max-bs + - "8" + - --max-running-requests + - "20" + - --tp + - "16" # Size of Tensor Parallelism + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20000 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + resources: + limits: + nvidia.com/gpu: "8" + volumeMounts: + - mountPath: /dev/shm + name: dshm + - name: model + mountPath: /work/models + - name: ib + mountPath: /dev/infiniband + volumes: + - name: dshm + emptyDir: + medium: Memory + - name: ib + hostPath: + path: /dev/infiniband + - name: model + hostPath: + path: /data1/models/deepseek_v3_moe +--- +apiVersion: v1 +kind: Service +metadata: + name: sglang-leader +spec: + selector: + leaderworkerset.sigs.k8s.io/name: sglang + role: leader + ports: + - protocol: TCP + port: 40000 + targetPort: 40000 + +``` + +* Then use `kubectl apply -f lws.yaml` you will get this output. + +```text +NAME READY STATUS RESTARTS AGE +sglang-0 0/1 Running 0 9s +sglang-0-1 1/1 Running 0 9s +``` + +Wait for the sglang leader (`sglang-0`) status to change to 1/1, which indicates it is `Ready`. + +You can use the command `kubectl logs -f sglang-0` to view the logs of the leader node. + +Once successful, you should see output like this: + +```text +[2025-02-17 05:27:24 TP1] Capture cuda graph end. Time elapsed: 84.89 s +[2025-02-17 05:27:24 TP6] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP0] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP7] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP3] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP2] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP4] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP1] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP5] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24] INFO: Started server process [1] +[2025-02-17 05:27:24] INFO: Waiting for application startup. +[2025-02-17 05:27:24] INFO: Application startup complete. +[2025-02-17 05:27:24] INFO: Uvicorn running on http://0.0.0.0:40000 (Press CTRL+C to quit) +[2025-02-17 05:27:25] INFO: 127.0.0.1:48908 - "GET /get_model_info HTTP/1.1" 200 OK +[2025-02-17 05:27:25 TP0] Prefill batch. #new-seq: 1, #new-token: 7, #cached-token: 0, cache hit rate: 0.00%, token usage: 0.00, #running-req: 0, #queue-req: 0 +[2025-02-17 05:27:32] INFO: 127.0.0.1:48924 - "POST /generate HTTP/1.1" 200 OK +[2025-02-17 05:27:32] The server is fired up and ready to roll! +``` + +If it doesn’t start up successfully, please follow these steps to check for any remaining issues. Thanks! + +### Debug + +* Set `NCCL_DEBUG=TRACE` to check if it is a NCCL communication problem. + +This should resolve most NCCL-related issues. + +***Notice: If you find that NCCL_DEBUG=TRACE is not effective in the container environment, but the process is stuck or you encounter hard-to-diagnose issues, try switching to a different container image. Some images may not handle standard error output properly.*** + +#### RoCE scenario + +* Please make sure that RDMA devices are available in the cluster environment. +* Please make sure that the nodes in the cluster have Mellanox NICs with RoCE. In this example, we use Mellanox ConnectX 5 model NICs, and the proper OFED driver has been installed. If not, please refer to the document [Install OFED Driver](https://docs.nvidia.com/networking/display/mlnxofedv461000/installing+mellanox+ofed) to install the driver. +* Check your env: + + ```shell + $ lspci -nn | grep Eth | grep Mellanox + 0000:7f:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0000:7f:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0000:c7:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0000:c7:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0001:08:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0001:08:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0001:a2:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0001:a2:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + ``` + +* Check the OFED driver: + + ```shell + ofed_info -s + OFED-internal-23.07-0.5.0: + ``` + +* Show RDMA link status and check IB devices: + + ```shell + $ rdma link show + 8/1: mlx5_bond_0/1: state ACTIVE physical_state LINK_UP netdev reth0 + 9/1: mlx5_bond_1/1: state ACTIVE physical_state LINK_UP netdev reth2 + 10/1: mlx5_bond_2/1: state ACTIVE physical_state LINK_UP netdev reth4 + 11/1: mlx5_bond_3/1: state ACTIVE physical_state LINK_UP netdev reth6 + + $ ibdev2netdev + 8/1: mlx5_bond_0/1: state ACTIVE physical_state LINK_UP netdev reth0 + 9/1: mlx5_bond_1/1: state ACTIVE physical_state LINK_UP netdev reth2 + 10/1: mlx5_bond_2/1: state ACTIVE physical_state LINK_UP netdev reth4 + 11/1: mlx5_bond_3/1: state ACTIVE physical_state LINK_UP netdev reth6 + ``` + +* Test RoCE network speed on the host: + + ```shell + yum install qperf + # for server： + execute qperf + # for client + qperf -t 60 -cm1 rc_rdma_write_bw + ``` + +* Check RDMA accessible in your container: + + ```shell + # ibv_devices + # ibv_devinfo + ``` + +## Keys to success + +* In the YAML configuration above, pay attention to the NCCL environment variable. For older versions of NCCL, you should check the NCCL_IB_GID_INDEX environment setting. +* NCCL_SOCKET_IFNAME is also crucial, but in a containerized environment, this typically isn’t an issue. +* In some cases, it’s necessary to configure GLOO_SOCKET_IFNAME correctly. +* NCCL_DEBUG is essential for troubleshooting, but I've found that sometimes it doesn't show error logs within containers. This could be related to the Docker image you're using. You may want to try switching images if needed. +* Avoid using Docker images based on Ubuntu 18.04, as they tend to have compatibility issues. + +## Remaining issues + +* In Kubernetes, Docker, or Containerd environments, we use hostNetwork to prevent performance degradation. +* We utilize privileged mode, which isn’t secure. Additionally, in containerized environments, full GPU isolation cannot be achieved. + +## TODO + +* Integrated with [k8s-rdma-shared-dev-plugin](https://github.com/Mellanox/k8s-rdma-shared-dev-plugin). diff --git a/docs/references/multi_node_deployment/lws_pd/lws-examples/d-svc.yaml b/docs/references/multi_node_deployment/lws_pd/lws-examples/d-svc.yaml new file mode 100644 index 000000000..27f98009e --- /dev/null +++ b/docs/references/multi_node_deployment/lws_pd/lws-examples/d-svc.yaml @@ -0,0 +1,12 @@ +apiVersion: v1 +kind: Service +metadata: + name: deepseekr10528-decode-main +spec: + selector: + leaderworkerset.sigs.k8s.io/name: deepseekr10528-decode-main + role: leader + ports: + - protocol: TCP + port: 30000 + targetPort: 30000 diff --git a/docs/references/multi_node_deployment/lws_pd/lws-examples/d.yaml b/docs/references/multi_node_deployment/lws_pd/lws-examples/d.yaml new file mode 100644 index 000000000..ac1d295eb --- /dev/null +++ b/docs/references/multi_node_deployment/lws_pd/lws-examples/d.yaml @@ -0,0 +1,290 @@ +apiVersion: leaderworkerset.x-k8s.io/v1 +kind: LeaderWorkerSet +metadata: + name: deepseekr10528-decode-main +spec: + leaderWorkerTemplate: + leaderTemplate: + metadata: + labels: + role: leader + spec: + containers: + - command: + - python3 + - -m + - sglang.launch_server + - --port + - "30000" + - --host + - "0.0.0.0" + - --model-path + - /work/models + - --chunked-prefill-size + - "262144" + - --page-size + - "64" + - --enable-dp-attention + - --enable-dp-lm-head + - --dp-size + - "16" + - --moe-a2a-backend + - deepep + - --disaggregation-mode + - decode + - --mem-fraction-static + - "0.849" + - --context-length + - "32768" + - --disaggregation-ib-device + - "mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3" + - --cuda-graph-max-bs + - "64" + - --max-running-requests + - "2048" + - --tp-size + - "16" # Size of Tensor Parallelism + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20102 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --ep-num-redundant-experts + - "32" + - --moe-dense-tp-size + - "1" + env: + - name: CUDA_LAUNCH_BLOCKING + value: "0" + - name: NVSHMEM_IB_GID_INDEX + value: "3" + - name: NVSHMEM_ENABLE_NIC_PE_MAPPING + value: "1" + - name: NVSHMEM_HCA_PE_MAPPING + value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" + - name: NCCL_IB_QPS_PER_CONNECTION + value: "8" + - name: NCCL_IB_SPLIT_DATA_ON_QPS + value: "1" + - name: NCCL_NET_PLUGIN + value: "none" + - name: NCCL_IB_TC + value: "136" + - name: NCCL_MIN_NCHANNELS + value: "4" + - name: NCCL_IB_SL + value: "5" + - name: MC_TE_METRIC + value: "true" + - name: SGLANG_MOONCAKE_TRANS_THREAD + value: "16" + - name: SGL_ENABLE_JIT_DEEPGEMM + value: "1" + - name: NCCL_IB_HCA + value: ^=mlx5_0,mlx5_5,mlx5_6 + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + image: lmsysorg/sglang:latest + name: sglang-leader + ports: + - containerPort: 30000 + protocol: TCP + readinessProbe: + periodSeconds: 30 + tcpSocket: + port: 30000 + resources: + limits: + nvidia.com/gpu: "8" + securityContext: + capabilities: + add: + - IPC_LOCK + privileged: true + volumeMounts: + - mountPath: /root/.cache + name: sgl-cache + - mountPath: /dev/shm + name: dshm + - mountPath: /work/models + name: model + - mountPath: /dev/infiniband + name: ib + - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs + name: cf + dnsPolicy: ClusterFirstWithHostNet + hostIPC: true + hostNetwork: true + nodeSelector: + # should modify according your deployment env + pd: "yes" + tolerations: + # should modify according your deployment env + - key: bopd + operator: Exists + - key: node-role + operator: Exists + volumes: + - hostPath: + path: /data1/sgl_cache1 + type: DirectoryOrCreate + name: sgl-cache + - emptyDir: + medium: Memory + name: dshm + - hostPath: + path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 + name: model + - hostPath: + path: /dev/infiniband + name: ib + - hostPath: + path: /data1/maas_hosted_models/models/fused_moe_triton/configs + name: cf + restartPolicy: RecreateGroupOnPodRestart + size: 2 + workerTemplate: + metadata: {} + spec: + containers: + - command: + - python3 + - -m + - sglang.launch_server + - --model-path + - /work/models + - --chunked-prefill-size + - "262144" + - --page-size + - "64" + - --enable-dp-attention + - --enable-dp-lm-head + - --dp-size + - "16" + - --moe-a2a-backend + - deepep + - --disaggregation-mode + - decode + - --mem-fraction-static + - "0.849" + - --context-length + - "32768" + - --disaggregation-ib-device + - "mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3" + - --cuda-graph-max-bs + - "64" + - --max-running-requests + - "2048" + - --tp-size + - "16" # Size of Tensor Parallelism + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20102 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --ep-num-redundant-experts + - "32" + - --moe-dense-tp-size + - "1" + env: + - name: NVSHMEM_IB_TRAFFIC_CLASS + value: "16" + - name: NVSHMEM_IB_GID_INDEX + value: "3" + - name: NVSHMEM_ENABLE_NIC_PE_MAPPING + value: "1" + - name: NVSHMEM_HCA_PE_MAPPING + value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" + - name: NCCL_IB_QPS_PER_CONNECTION + value: "8" + - name: NCCL_IB_SPLIT_DATA_ON_QPS + value: "1" + - name: NCCL_NET_PLUGIN + value: "none" + - name: NCCL_IB_TC + value: "136" + - name: NCCL_MIN_NCHANNELS + value: "4" + - name: MC_TE_METRIC + value: "true" + - name: NCCL_IB_SL + value: "5" + - name: SGLANG_MOONCAKE_TRANS_THREAD + value: "16" + - name: SGL_ENABLE_JIT_DEEPGEMM + value: "1" + - name: NCCL_IB_HCA + value: ^=mlx5_0,mlx5_5,mlx5_6 + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + image: lmsysorg/sglang:latest + name: sglang-worker + ports: + - containerPort: 30001 + resources: + limits: + nvidia.com/gpu: "8" + securityContext: + capabilities: + add: + - IPC_LOCK + privileged: true + volumeMounts: + - mountPath: /root/.cache + name: sgl-cache + - mountPath: /dev/shm + name: dshm + - mountPath: /work/models + name: model + - mountPath: /dev/infiniband + name: ib + - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs + name: cf + dnsPolicy: ClusterFirstWithHostNet + hostIPC: true + hostNetwork: true + nodeSelector: + # should modify according your deployment env + pd: "yes" + tolerations: + # should modify according your deployment env + - key: bopd + operator: Exists + - key: node-role + operator: Exists + volumes: + - hostPath: + path: /data1/sgl_cache1 + type: DirectoryOrCreate + name: sgl-cache + - emptyDir: + medium: Memory + name: dshm + - hostPath: + path: /dev/infiniband + name: ib + - hostPath: + # modify according to you deployment env + path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 + name: model + - hostPath: + # modify according to you deployment env + path: /data1/maas_hosted_models/models/fused_moe_triton/configs + name: cf + networkConfig: + subdomainPolicy: Shared + replicas: 1 + rolloutStrategy: + rollingUpdateConfiguration: + maxSurge: 0 + maxUnavailable: 1 + type: RollingUpdate + startupPolicy: LeaderCreated diff --git a/docs/references/multi_node_deployment/lws_pd/lws-examples/lb.yaml b/docs/references/multi_node_deployment/lws_pd/lws-examples/lb.yaml new file mode 100644 index 000000000..4ca690969 --- /dev/null +++ b/docs/references/multi_node_deployment/lws_pd/lws-examples/lb.yaml @@ -0,0 +1,56 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: deepseekr10528-lb-main + labels: + app: deepseekr10528-lb +spec: + replicas: 1 + selector: + matchLabels: + app: deepseekr10528-lb + template: + metadata: + labels: + app: deepseekr10528-lb + spec: + nodeSelector: + bo: "yes" + tolerations: + - key: bopd + operator: Exists + - key: node-role + operator: Exists + containers: + - name: sgl-minilb + image: lmsysorg/sglang:latest + command: + - python + - -m + - sglang_router.launch_router + - --pd-disaggregation + - --prefill + - http://deepseekr10528-prefill-main:30000 + - --decode + - http://deepseekr10528-decode-main:30000 + - --host + - 0.0.0.0 + - --port + - "8000" + ports: + - containerPort: 8000 + +--- +apiVersion: v1 +kind: Service +metadata: + name: deepseekr10528-lb-service +spec: + type: NodePort # NodePort is easy to test, you can also specify `ClusterIP` + selector: + app: deepseekr10528-lb + ports: + - protocol: TCP + port: 8000 # Service Port（In-Cluster） + targetPort: 8000 # Exposed Container + nodePort: 30800 diff --git a/docs/references/multi_node_deployment/lws_pd/lws-examples/p-svc.yaml b/docs/references/multi_node_deployment/lws_pd/lws-examples/p-svc.yaml new file mode 100644 index 000000000..6826a13df --- /dev/null +++ b/docs/references/multi_node_deployment/lws_pd/lws-examples/p-svc.yaml @@ -0,0 +1,12 @@ +apiVersion: v1 +kind: Service +metadata: + name: deepseekr10528-prefill-main +spec: + selector: + leaderworkerset.sigs.k8s.io/name: deepseekr10528-prefill-main + role: leader + ports: + - protocol: TCP + port: 30000 + targetPort: 30000 diff --git a/docs/references/multi_node_deployment/lws_pd/lws-examples/p.yaml b/docs/references/multi_node_deployment/lws_pd/lws-examples/p.yaml new file mode 100644 index 000000000..62df262bb --- /dev/null +++ b/docs/references/multi_node_deployment/lws_pd/lws-examples/p.yaml @@ -0,0 +1,304 @@ +apiVersion: leaderworkerset.x-k8s.io/v1 +kind: LeaderWorkerSet +metadata: + name: deepseekr10528-prefill-main +spec: + leaderWorkerTemplate: + leaderTemplate: + metadata: + labels: + role: leader + spec: + containers: + - command: + - python3 + - -m + - sglang.launch_server + - --port + - "30000" + - --host + - "0.0.0.0" + - --model-path + - /work/models + - --disaggregation-ib-device + # should modify according your rdma env + - mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 + - --chunked-prefill-size + - "524288" + - --max-prefill-tokens + - "32768" + - --page-size + - "64" + - --ep-dispatch-algorithm + - dynamic + - --eplb-algorithm + - deepseek + - --enable-dp-lm-head + - --enable-dp-attention + - --dp-size + - "16" + - --disable-radix-cache + - --moe-a2a-backend + - deepep + - --disaggregation-mode + - prefill + - --mem-fraction-static + - "0.7" + - --context-length + - "32768" + - --tp + - "16" + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20102 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --ep-num-redundant-experts + - "32" + - --moe-dense-tp-size + - "1" + - --max-running-requests + - "1024" + env: + - name: NVSHMEM_HCA_PE_MAPPING + # should modify according your rdma env + value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" + - name: NVSHMEM_IB_GID_INDEX + value: "3" + - name: NVSHMEM_ENABLE_NIC_PE_MAPPING + value: "1" + - name: SGLANG_SET_CPU_AFFINITY + value: "true" + - name: SGL_ENABLE_JIT_DEEPGEMM + value: "1" + - name: NCCL_IB_QPS_PER_CONNECTION + value: "8" + - name: NCCL_IB_SPLIT_DATA_ON_QPS + value: "1" + - name: NCCL_NET_PLUGIN + value: none + - name: NCCL_IB_TC + value: "136" + - name: NCCL_MIN_NCHANNELS + value: "4" + - name: MC_TE_METRIC + value: "false" + - name: NCCL_IB_SL + value: "5" + - name: NCCL_IB_HCA + value: ^=mlx5_0,mlx5_5,mlx5_6 + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + image: lmsysorg/sglang:latest + name: sglang-leader + ports: + - containerPort: 30000 + protocol: TCP + readinessProbe: + periodSeconds: 30 + tcpSocket: + port: 30000 + resources: + limits: + nvidia.com/gpu: "8" + securityContext: + capabilities: + add: + - IPC_LOCK + privileged: true + volumeMounts: + - mountPath: /dev/shm + name: dshm + - mountPath: /work/models + name: model + - mountPath: /dev/infiniband + name: ib + - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs + name: cf + - mountPath: /root/.cache + name: sgl-cache + dnsPolicy: ClusterFirstWithHostNet + hostIPC: true + hostNetwork: true + nodeSelector: + # should modify according your deployment env + pd: "yes" + tolerations: + # should modify according your deployment env + - key: bopd + operator: Exists + - key: node-role + operator: Exists + volumes: + - emptyDir: + medium: Memory + name: dshm + - hostPath: + path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 + name: model + - hostPath: + path: /dev/infiniband + name: ib + - hostPath: + path: /data1/maas_hosted_models/models/fused_moe_triton/configs + name: cf + - hostPath: + path: /data1/sgl_cache + type: DirectoryOrCreate + name: sgl-cache + restartPolicy: RecreateGroupOnPodRestart + size: 2 + workerTemplate: + metadata: {} + spec: + containers: + - command: + - python3 + - -m + - sglang.launch_server + - --model-path + - /work/models + - --disaggregation-ib-device + # should modify according your rdma env + - mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 + - --chunked-prefill-size + - "524288" + - --max-prefill-tokens + - "32768" + - --page-size + - "64" + - --ep-dispatch-algorithm + - dynamic + - --eplb-algorithm + - deepseek + # - --deepep-config + # - /home/aiges/tuned/tuned_8sms.json + # can be tuned using deepep test scripts + - --enable-dp-lm-head + - --enable-dp-attention + - --dp-size + - "16" + - --disable-radix-cache + - --moe-a2a-backend + - deepep + - --disaggregation-mode + - prefill + - --mem-fraction-static + - "0.7" + - --context-length + - "32768" + - --tp + - "16" + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20102 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --ep-num-redundant-experts + - "32" + - --moe-dense-tp-size + - "1" + - --max-running-requests + - "1024" + env: + - name: SGLANG_SET_CPU_AFFINITY + value: "true" + - name: NVSHMEM_HCA_PE_MAPPING + # should modify according your rdma env + value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" + - name: NCCL_IB_HCA + value: ^=mlx5_0,mlx5_5,mlx5_6 + - name: NVSHMEM_IB_TRAFFIC_CLASS + value: "16" + - name: NVSHMEM_IB_GID_INDEX + value: "3" + - name: NVSHMEM_ENABLE_NIC_PE_MAPPING + value: "1" + - name: CUDA_LAUNCH_BLOCKING + value: "0" + - name: SGLANG_MOONCAKE_TRANS_THREAD + value: "8" + - name: SGL_ENABLE_JIT_DEEPGEMM + value: "1" + - name: SGL_CHUNKED_PREFIX_CACHE_THRESHOLD + value: "0" + - name: NCCL_IB_QPS_PER_CONNECTION + value: "8" + - name: NCCL_IB_SPLIT_DATA_ON_QPS + value: "1" + - name: NCCL_NET_PLUGIN + value: none + - name: NCCL_IB_TC + value: "136" + - name: NCCL_MIN_NCHANNELS + value: "4" + - name: MC_TE_METRIC + value: "true" + - name: NCCL_IB_SL + value: "5" + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + image: lmsysorg/sglang:latest + name: sglang-worker + ports: + - containerPort: 30001 + protocol: TCP + resources: + limits: + nvidia.com/gpu: "8" + securityContext: + capabilities: + add: + - IPC_LOCK + privileged: true + volumeMounts: + - mountPath: /root/.cache + name: sgl-cache + - mountPath: /dev/shm + name: dshm + - mountPath: /work/models + name: model + - mountPath: /dev/infiniband + name: ib + - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs + name: cf + dnsPolicy: ClusterFirstWithHostNet + hostIPC: true + hostNetwork: true + nodeSelector: + # should modify according your deployment env + pd: "yes" + tolerations: + # should modify according your deployment env + - key: bopd + operator: Exists + - key: node-role + operator: Exists + volumes: + - emptyDir: + medium: Memory + name: dshm + - hostPath: + path: /dev/infiniband + name: ib + - hostPath: + # modify according to you deployment env + path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 + name: model + - hostPath: + # modify according to you deployment env + path: /data1/maas_hosted_models/models/fused_moe_triton/configs + name: cf + - hostPath: + # modify according to you deployment env + path: /data1/sgl_cache + type: DirectoryOrCreate + name: sgl-cache diff --git a/docs/references/multi_node_deployment/lws_pd/lws_pd_deploy.md b/docs/references/multi_node_deployment/lws_pd/lws_pd_deploy.md new file mode 100644 index 000000000..eb8454997 --- /dev/null +++ b/docs/references/multi_node_deployment/lws_pd/lws_pd_deploy.md @@ -0,0 +1,783 @@ +# LWS Based PD Deploy + +## 0. Prerequisites + +1. k8s >=1.26 +2. lws installed on k8s. + +## 1. Image Preparation + +`lmsysorg/sglang:deepep` + +## 2. Deployment Manifest Files + +***Notice: We will package all deployment files into Helm Chart format in the near future. Interested community members can contact us to contribute*** + +### Prefill + +Prefill manifest file [prefill.yaml](lws-examples/p.yaml) + +*Note: The NodeSelector section, model location section, and taint toleration section can be adjusted according to your actual deployment environment* + +```yaml +apiVersion: leaderworkerset.x-k8s.io/v1 +kind: LeaderWorkerSet +metadata: + name: deepseekr10528-prefill-main +spec: + leaderWorkerTemplate: + leaderTemplate: + metadata: + labels: + role: leader + spec: + containers: + - command: + - python3 + - -m + - sglang.launch_server + - --port + - "30000" + - --host + - "0.0.0.0" + - --model-path + - /work/models + - --disaggregation-ib-device + # should modify according your rdma env + - mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 + - --chunked-prefill-size + - "524288" + - --max-prefill-tokens + - "32768" + - --page-size + - "64" + # - --init-expert-location + # - /home/aiges/tuned/attachment_ep_statistics/prefill_in1024.json + - --ep-dispatch-algorithm + - dynamic + - --eplb-algorithm + - deepseek + # - --deepep-config + # - /home/aiges/tuned/tuned_8sms.json + - --enable-dp-lm-head + - --enable-dp-attention + - --dp-size + - "16" + - --disable-radix-cache + - --moe-a2a-backend + - deepep + - --disaggregation-mode + - prefill + - --mem-fraction-static + - "0.7" + - --context-length + - "32768" + - --tp + - "16" + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20102 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --ep-num-redundant-experts + - "32" + - --moe-dense-tp-size + - "1" + - --max-running-requests + - "1024" + env: +# - name: NVSHMEM_HCA_PE_MAPPING +# value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" +# - name: NVSHMEM_HCA_LIST +# value: "mlx5_bond_0:1,mlx5_bond_1:1,mlx5_bond_2:1,mlx5_bond_3:1" + - name: NVSHMEM_IB_GID_INDEX + value: "3" + - name: NVSHMEM_ENABLE_NIC_PE_MAPPING + value: "1" + - name: SGLANG_SET_CPU_AFFINITY + value: "true" + - name: SGL_ENABLE_JIT_DEEPGEMM + value: "1" + - name: NCCL_IB_QPS_PER_CONNECTION + value: "8" + - name: NCCL_IB_SPLIT_DATA_ON_QPS + value: "1" + - name: NCCL_NET_PLUGIN + value: none + - name: NCCL_IB_TC + value: "136" + - name: NCCL_MIN_NCHANNELS + value: "4" + - name: MC_TE_METRIC + value: "false" + - name: NCCL_IB_SL + value: "5" + - name: NCCL_IB_HCA + value: ^=mlx5_0,mlx5_5,mlx5_6 + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + image: lmsysorg/sglang:deepep + name: sglang-leader + ports: + - containerPort: 30000 + protocol: TCP + readinessProbe: + periodSeconds: 30 + tcpSocket: + port: 30000 + resources: + limits: + nvidia.com/gpu: "8" + securityContext: + capabilities: + add: + - IPC_LOCK + privileged: true + volumeMounts: + - mountPath: /dev/shm + name: dshm + - mountPath: /work/models + name: model + - mountPath: /dev/infiniband + name: ib + - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs + name: cf + - mountPath: /root/.cache + name: sgl-cache + dnsPolicy: ClusterFirstWithHostNet + hostIPC: true + hostNetwork: true + nodeSelector: + pd: "yes" + tolerations: + - key: pd + operator: Exists + - key: node-role + operator: Exists + volumes: + - emptyDir: + medium: Memory + name: dshm + - hostPath: + # modify according to you deployment env + path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 + name: model + - hostPath: + path: /dev/infiniband + name: ib + - hostPath: + # modify according to you deployment env + path: /data1/maas_hosted_models/models/fused_moe_triton/configs + name: cf + - hostPath: + # modify according to you deployment env + path: /data1/sgl_cache + type: DirectoryOrCreate + name: sgl-cache + restartPolicy: RecreateGroupOnPodRestart + size: 2 + workerTemplate: + metadata: {} + spec: + containers: + - command: + - python3 + - -m + - sglang.launch_server + - --model-path + - /work/models + - --disaggregation-ib-device + - mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 + - --chunked-prefill-size + - "524288" + - --max-prefill-tokens + - "32768" + - --page-size + - "64" + #- --init-expert-location + #- /home/aiges/tuned/attachment_ep_statistics/prefill_in1024.json + - --ep-dispatch-algorithm + - dynamic + - --eplb-algorithm + - deepseek +# - --deepep-config +# - /home/aiges/tuned/tuned_8sms.json + - --enable-dp-lm-head + - --enable-dp-attention + - --dp-size + - "16" + - --disable-radix-cache + - --moe-a2a-backend + - deepep + - --disaggregation-mode + - prefill + - --mem-fraction-static + - "0.7" + - --context-length + - "32768" + - --tp + - "16" + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20102 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --ep-num-redundant-experts + - "32" + - --moe-dense-tp-size + - "1" + - --max-running-requests + - "1024" + env: + - name: SGLANG_SET_CPU_AFFINITY + value: "true" + - name: SGLANG_HACK_DEEPEP_NUM_SMS + value: "8" + - name: SGLANG_HACK_DEEPEP_NEW_MODE + value: "0" +# - name: NVSHMEM_HCA_PE_MAPPING +# value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" +# - name: NVSHMEM_HCA_LIST +# value: "mlx5_bond_0:1,mlx5_bond_1:1,mlx5_bond_2:1,mlx5_bond_3:1" + - name: NCCL_IB_HCA + value: ^=mlx5_0,mlx5_5,mlx5_6 + - name: NVSHMEM_IB_TRAFFIC_CLASS + value: "16" + - name: NVSHMEM_IB_GID_INDEX + value: "3" + - name: NVSHMEM_ENABLE_NIC_PE_MAPPING + value: "1" + - name: CUDA_LAUNCH_BLOCKING + value: "0" + - name: SGLANG_MOONCAKE_TRANS_THREAD + value: "8" + - name: SGL_ENABLE_JIT_DEEPGEMM + value: "1" + - name: SGL_CHUNKED_PREFIX_CACHE_THRESHOLD + value: "0" + - name: NCCL_IB_QPS_PER_CONNECTION + value: "8" + - name: NCCL_IB_SPLIT_DATA_ON_QPS + value: "1" + - name: NCCL_NET_PLUGIN + value: none + - name: NCCL_IB_TC + value: "136" + - name: NCCL_MIN_NCHANNELS + value: "4" + - name: MC_TE_METRIC + value: "true" + - name: NCCL_IB_SL + value: "5" + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + image: lmsysorg/sglang:deepep + name: sglang-worker + ports: + - containerPort: 30001 + protocol: TCP + resources: + limits: + nvidia.com/gpu: "8" + securityContext: + capabilities: + add: + - IPC_LOCK + privileged: true + volumeMounts: + + - mountPath: /root/.cache + name: sgl-cache + - mountPath: /dev/shm + name: dshm + - mountPath: /work/models + name: model + - mountPath: /dev/infiniband + name: ib + - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs + name: cf + dnsPolicy: ClusterFirstWithHostNet + hostIPC: true + hostNetwork: true + nodeSelector: + pd: "yes" + tolerations: + - key: pd + operator: Exists + - key: node-role + operator: Exists + volumes: + - emptyDir: + medium: Memory + name: dshm + - hostPath: + path: /dev/infiniband + name: ib + - hostPath: + path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 + name: model + - hostPath: + path: /data1/maas_hosted_models/models/fused_moe_triton/configs + name: cf + - hostPath: + path: /data1/sgl_cache + type: DirectoryOrCreate + name: sgl-cache + +``` + +### Decode + +Decode node deployment manifest file [decode.yaml](lws-examples/d.yaml) + +*Note: The NodeSelector section, model location section, and taint toleration section can be adjusted according to your actual deployment environment* + +```yaml +apiVersion: leaderworkerset.x-k8s.io/v1 +kind: LeaderWorkerSet +metadata: + name: deepseekr10528-decode-main +spec: + leaderWorkerTemplate: + leaderTemplate: + metadata: + labels: + role: leader + spec: + containers: + - command: + - python3 + - -m + - sglang.launch_server + - --port + - "30000" + - --host + - "0.0.0.0" + - --model-path + - /work/models + - --chunked-prefill-size + - "262144" + - --page-size + - "64" + - --enable-dp-attention + - --enable-dp-lm-head + - --dp-size + - "16" + - --moe-a2a-backend + - deepep + - --disaggregation-mode + - decode + - --mem-fraction-static + - "0.849" + - --context-length + - "32768" + - --disaggregation-ib-device + - "mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3" + - --cuda-graph-max-bs + - "64" + - --max-running-requests + - "2048" + - --tp-size + - "16" # Size of Tensor Parallelism + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20102 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --ep-num-redundant-experts + - "32" + - --moe-dense-tp-size + - "1" + env: + - name: CUDA_LAUNCH_BLOCKING + value: "0" + - name: NVSHMEM_IB_GID_INDEX + value: "3" + - name: NVSHMEM_ENABLE_NIC_PE_MAPPING + value: "1" + - name: NCCL_IB_QPS_PER_CONNECTION + value: "8" + - name: NCCL_IB_SPLIT_DATA_ON_QPS + value: "1" + - name: NCCL_NET_PLUGIN + value: "none" + - name: NCCL_IB_TC + value: "136" + - name: NCCL_MIN_NCHANNELS + value: "4" + - name: NCCL_IB_SL + value: "5" + - name: MC_TE_METRIC + value: "true" + - name: SGLANG_MOONCAKE_TRANS_THREAD + value: "16" + - name: SGL_ENABLE_JIT_DEEPGEMM + value: "1" + - name: NCCL_IB_HCA + value: ^=mlx5_0,mlx5_5,mlx5_6 + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + image: lmsysorg/sglang:deepep + name: sglang-leader + ports: + - containerPort: 30000 + protocol: TCP + readinessProbe: + periodSeconds: 30 + tcpSocket: + port: 30000 + resources: + limits: + nvidia.com/gpu: "8" + securityContext: + capabilities: + add: + - IPC_LOCK + privileged: true + volumeMounts: + - mountPath: /root/.cache + name: sgl-cache + - mountPath: /dev/shm + name: dshm + - mountPath: /work/models + name: model + - mountPath: /dev/infiniband + name: ib + - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs + name: cf + dnsPolicy: ClusterFirstWithHostNet + hostIPC: true + hostNetwork: true + nodeSelector: + pd: "yes" + tolerations: + - key: pd + operator: Exists + - key: node-role + operator: Exists + volumes: + - hostPath: + path: /data1/sgl_cache1 + type: DirectoryOrCreate + name: sgl-cache + - emptyDir: + medium: Memory + name: dshm + - hostPath: + path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 + name: model + - hostPath: + path: /dev/infiniband + name: ib + - hostPath: + path: /data1/maas_hosted_models/models/fused_moe_triton/configs + name: cf + restartPolicy: RecreateGroupOnPodRestart + size: 2 + workerTemplate: + metadata: {} + spec: + containers: + - command: + - python3 + - -m + - sglang.launch_server + - --model-path + - /work/models + - --chunked-prefill-size + - "262144" + - --page-size + - "64" + - --enable-dp-attention + - --enable-dp-lm-head + #- --enable-two-batch-overlap + - --dp-size + - "16" + - --moe-a2a-backend + - deepep + - --disaggregation-mode + - decode + - --mem-fraction-static + - "0.849" + - --context-length + - "32768" + - --disaggregation-ib-device + # should modify according your rdma env + - "mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3" + - --cuda-graph-max-bs + - "64" + - --max-running-requests + - "2048" + - --tp-size + - "16" # Size of Tensor Parallelism + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20102 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --ep-num-redundant-experts + - "32" + - --moe-dense-tp-size + - "1" + env: + - name: SGLANG_HACK_DEEPEP_NUM_SMS + value: "24" + - name: SGLANG_HACK_DEEPEP_NEW_MODE + value: "0" + - name: NVSHMEM_IB_TRAFFIC_CLASS + value: "16" + - name: NVSHMEM_IB_GID_INDEX + value: "3" + - name: NVSHMEM_ENABLE_NIC_PE_MAPPING + value: "1" + - name: NCCL_IB_QPS_PER_CONNECTION + value: "8" + - name: NCCL_IB_SPLIT_DATA_ON_QPS + value: "1" + - name: NCCL_NET_PLUGIN + value: "none" + - name: NCCL_IB_TC + value: "136" + - name: NCCL_MIN_NCHANNELS + value: "4" + - name: MC_TE_METRIC + value: "true" + - name: NCCL_IB_SL + value: "5" + - name: SGLANG_MOONCAKE_TRANS_THREAD + value: "16" + - name: SGL_ENABLE_JIT_DEEPGEMM + value: "1" + - name: NCCL_IB_HCA + value: ^=mlx5_0,mlx5_5,mlx5_6 + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + image: lmsysorg/sglang:deepep + name: sglang-worker + ports: + - containerPort: 30001 + resources: + limits: + nvidia.com/gpu: "8" + securityContext: + capabilities: + add: + - IPC_LOCK + privileged: true + volumeMounts: + - mountPath: /root/.cache + name: sgl-cache + - mountPath: /dev/shm + name: dshm + - mountPath: /work/models + name: model + - mountPath: /dev/infiniband + name: ib + - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs + name: cf + dnsPolicy: ClusterFirstWithHostNet + hostIPC: true + hostNetwork: true + nodeSelector: + pd: "yes" + tolerations: + - key: pd + operator: Exists + - key: node-role + operator: Exists + volumes: + - hostPath: + path: /data1/sgl_cache1 + type: DirectoryOrCreate + name: sgl-cache + - emptyDir: + medium: Memory + name: dshm + - hostPath: + path: /dev/infiniband + name: ib + - hostPath: + # modify according to you deployment env + path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 + name: model + - hostPath: + # modify according to you deployment env + path: /data1/maas_hosted_models/models/fused_moe_triton/configs + name: cf + networkConfig: + subdomainPolicy: Shared + replicas: 1 + rolloutStrategy: + rollingUpdateConfiguration: + maxSurge: 0 + maxUnavailable: 1 + type: RollingUpdate + startupPolicy: LeaderCreated +``` + +Execute separately: + +```bash +kubectl apply -f p.yaml +kubectl apply -f d.yaml +``` + +At this point, we have completed the deployment of the 1P1D SGlang engine part. + +To allow our users to directly experience the model API, we still need a load balancer to handle sequential calls between prefill and decode. Different companies implement LBs differently, and the community will also officially release a new LB component written in Rust in the near future. + +Currently, we use a static K8S service + minilb approach to implement model API calls. + +### Creating Service for Prefill and Decode + +#### Create prefill k8s service +[p-svc.yaml](lws-examples/p-svc.yaml) +```yaml +apiVersion: v1 +kind: Service +metadata: + name: deepseekr10528-prefill-main +spec: + selector: + leaderworkerset.sigs.k8s.io/name: deepseekr10528-prefill-main + role: leader + ports: + - protocol: TCP + port: 30000 + targetPort: 30000 +``` +Execute `kubectl apply -f p-svc.yaml` + +#### Create decode k8s service +[d-svc.yaml](lws-examples/d-svc.yaml) +```yaml +apiVersion: v1 +kind: Service +metadata: + name: deepseekr10528-decode-main +spec: + selector: + leaderworkerset.sigs.k8s.io/name: deepseekr10528-decode-main + role: leader + ports: + - protocol: TCP + port: 30000 + targetPort: 30000 +``` +Execute `kubectl apply -f d-svc.yaml` + +#### Deploy minilb and lb service +[lb.yaml](lws-examples/lb.yaml) +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: deepseekr10528-lb-main + labels: + app: deepseekr10528-lb +spec: + replicas: 1 + selector: + matchLabels: + app: deepseekr10528-lb + template: + metadata: + labels: + app: deepseekr10528-lb + spec: + nodeSelector: + pd: "yes" + tolerations: + - key: pd + operator: Exists + - key: node-role + operator: Exists + containers: + - name: sgl-minilb + image: lmsysorg/sglang:deepep + command: + - python + - -m + - sglang_router.launch_router + - --pd-disaggregation + - --prefill + - http://deepseekr10528-prefill-main:30000 + - --decode + - http://deepseekr10528-decode-main:30000 + - --host + - 0.0.0.0 + - --port + - "8000" + ports: + - containerPort: 8000 +--- +apiVersion: v1 +kind: Service +metadata: + name: deepseekr10528-lb-service +spec: + type: NodePort + selector: + app: deepseekr10528-lb + ports: + - protocol: TCP + port: 8000 # Service Port（In-Cluster） + targetPort: 8000 # Exposed Container + nodePort: 30800 +``` +Execute `kubectl apply -f lb.yaml` + +After waiting for all model deployments to succeed, you will get the following output: + +```bash +[root@ecs-001]# kubectl get po +deepseekr10528-decode-main-0 1/1 Running 0 74m +deepseekr10528-decode-main-0-1 1/1 Running 0 74m +deepseekr10528-lb-main-9c5dbfc57-6lcbd 1/1 Running 0 22m +deepseekr10528-prefill-main-0 1/1 Running 0 74m +deepseekr10528-prefill-main-0-1 1/1 Running 0 74m +[root@ecs-cbm-x1-pd-cpu-001 main_doc]# kubectl get svc |grep dee +deepseekr10528-decode-main ClusterIP None 97m +deepseekr10528-lb-service NodePort 172.16.242.169 8000:30800/TCP 22m +deepseekr10528-prefill-main ClusterIP None 97m +``` + +At this point, select a nodePort:30800 to access: + +```bash +[root@ecs-001]# curl -X POST "http://{nodePort}:30800/v1/chat/completions" \ +> -H "Content-Type: application/json" \ +> -H "Authorization: Bearer None" \ +> -d '{ +> "rid":"ccccdd", +> "model": "r1", +> "messages": [ +> {"role": "system", "content": "0: You are a helpful AI assistant"}, +> {"role": "user", "content": "你是谁？."} +> ], +> "max_tokens":221 +> }' +{"id":"ccccdd","object":"chat.completion","created":1750252498,"model":"qwen2","choices":[{"index":0,"message":{"role":"assistant","content":"\n嗯，用户问了一个很基础的自我介绍问题"你是谁？"。这可能是第一次互动时的常规开场白，也可能是想确认我的身份和功能范围。\n\n用户没有提供任何背景信息，语气简洁中性。这种场景下新用户的可能性较高，需要给出清晰友好的自我介绍，同时突出实用价值来降低陌生感。\n\n考虑到中文用户，应该用简体中文回复。重点要说明三点：身份归属（深度求索）、功能定位（AI助手）、服务范围（学习/工作/生活）。结尾用开放性问题引导对话很关键——既能了解需求，又能避免让用户面对空白输入框时不知所措。\n\n用波浪线结尾可以软化语气，那个笑脸表情😊刚好能中和AI的机械感。不过要控制表情符号数量，避免显得轻浮。\n\n你好呀！我是你的AI助手，由深度求索公司（DeepSeek）开发的语言模型，名字叫 **DeepSeek-R1**。你可以把我当成一个知识丰富、随叫随到的小帮手～😊\n\n我的任务就是陪你聊天、解答问题、","reasoning_content":null,"tool_calls":null},"logprobs":null,"finish_reason":"length","matched_stop":null}],"usage":{"prompt_tokens":14,"total_tokens":235,"completion_tokens":221,"prompt_tokens_details":null}} + +``` +## FAQ + +1. The current deployment startup parameters may not be fully compatible with all RDMA scenarios. Different RDMA NCCL-related environment configurations may be needed in different network environments. + +2. Some preset, optimized configurations for EPLB are not used here. You can adjust them according to [6017](https://github.com/sgl-project/sglang/issues/6017) as needed. diff --git a/docs/references/multi_node_deployment/multi_node.md b/docs/references/multi_node_deployment/multi_node.md new file mode 100644 index 000000000..204b60586 --- /dev/null +++ b/docs/references/multi_node_deployment/multi_node.md @@ -0,0 +1,90 @@ +# Multi-Node Deployment + +## Llama 3.1 405B + +**Run 405B (fp16) on Two Nodes** + +```bash +# replace 172.16.4.52:20000 with your own node ip address and port of the first node + +python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-405B-Instruct --tp 16 --dist-init-addr 172.16.4.52:20000 --nnodes 2 --node-rank 0 + +python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-405B-Instruct --tp 16 --dist-init-addr 172.16.4.52:20000 --nnodes 2 --node-rank 1 +``` + +Note that LLama 405B (fp8) can also be launched on a single node. + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-405B-Instruct-FP8 --tp 8 +``` + +## DeepSeek V3/R1 + +Please refer to [DeepSeek documents for reference](https://docs.sglang.ai/basic_usage/deepseek.html#running-examples-on-multi-node). + +## Multi-Node Inference on SLURM + +This example showcases how to serve SGLang server across multiple nodes by SLURM. Submit the following job to the SLURM cluster. + +``` +#!/bin/bash -l + +#SBATCH -o SLURM_Logs/%x_%j_master.out +#SBATCH -e SLURM_Logs/%x_%j_master.err +#SBATCH -D ./ +#SBATCH -J Llama-405B-Online-Inference-TP16-SGL + +#SBATCH --nodes=2 +#SBATCH --ntasks=2 +#SBATCH --ntasks-per-node=1 # Ensure 1 task per node +#SBATCH --cpus-per-task=18 +#SBATCH --mem=224GB +#SBATCH --partition="lmsys.org" +#SBATCH --gres=gpu:8 +#SBATCH --time=12:00:00 + +echo "[INFO] Activating environment on node $SLURM_PROCID" +if ! source ENV_FOLDER/bin/activate; then + echo "[ERROR] Failed to activate environment" >&2 + exit 1 +fi + +# Define parameters +model=MODEL_PATH +tp_size=16 + +echo "[INFO] Running inference" +echo "[INFO] Model: $model" +echo "[INFO] TP Size: $tp_size" + +# Set NCCL initialization address using the hostname of the head node +HEAD_NODE=$(scontrol show hostname "$SLURM_NODELIST" | head -n 1) +NCCL_INIT_ADDR="${HEAD_NODE}:8000" +echo "[INFO] NCCL_INIT_ADDR: $NCCL_INIT_ADDR" + +# Launch the model server on each node using SLURM +srun --ntasks=2 --nodes=2 --output="SLURM_Logs/%x_%j_node$SLURM_NODEID.out" \ + --error="SLURM_Logs/%x_%j_node$SLURM_NODEID.err" \ + python3 -m sglang.launch_server \ + --model-path "$model" \ + --grammar-backend "xgrammar" \ + --tp "$tp_size" \ + --dist-init-addr "$NCCL_INIT_ADDR" \ + --nnodes 2 \ + --node-rank "$SLURM_NODEID" & + +# Wait for the NCCL server to be ready on port 30000 +while ! nc -z "$HEAD_NODE" 30000; do + sleep 1 + echo "[INFO] Waiting for $HEAD_NODE:30000 to accept connections" +done + +echo "[INFO] $HEAD_NODE:30000 is ready to accept connections" + +# Keep the script running until the SLURM job times out +wait +``` + +Then, you can test the server by sending requests following other [documents](https://docs.sglang.ai/backend/openai_api_completions.html). + +Thanks for [aflah02](https://github.com/aflah02) for providing the example, based on his [blog post](https://aflah02.substack.com/p/multi-node-llm-inference-with-sglang). diff --git a/docs/references/multi_node_deployment/multi_node_index.rst b/docs/references/multi_node_deployment/multi_node_index.rst new file mode 100644 index 000000000..03411f5be --- /dev/null +++ b/docs/references/multi_node_deployment/multi_node_index.rst @@ -0,0 +1,13 @@ +Multi-Node Deployment +===================== + +.. toctree:: + :maxdepth: 1 + :caption: Multi-Node Deployment + + multi_node.md + deploy_on_k8s.md + lws_pd/lws_pd_deploy.md + +- `Deploying DeepSeek with PD Disaggregation and Large-Scale Expert Parallelism on 96 H100 GPUs `_ +- `Deploying Kimi K2 with PD Disaggregation and Large-Scale Expert Parallelism on 128 H200 GPUs `_ diff --git a/docs/references/production_metrics.md b/docs/references/production_metrics.md new file mode 100644 index 000000000..16afaca67 --- /dev/null +++ b/docs/references/production_metrics.md @@ -0,0 +1,217 @@ +# Production Metrics + +SGLang exposes the following metrics via Prometheus. You can enable it by adding `--enable-metrics` when you launch the server. + +An example of the monitoring dashboard is available in [examples/monitoring/grafana.json](https://github.com/sgl-project/sglang/blob/main/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json). + +Here is an example of the metrics: + +``` +$ curl http://localhost:30000/metrics +# HELP sglang:prompt_tokens_total Number of prefill tokens processed. +# TYPE sglang:prompt_tokens_total counter +sglang:prompt_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct"} 8.128902e+06 +# HELP sglang:generation_tokens_total Number of generation tokens processed. +# TYPE sglang:generation_tokens_total counter +sglang:generation_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.557572e+06 +# HELP sglang:token_usage The token usage +# TYPE sglang:token_usage gauge +sglang:token_usage{model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.28 +# HELP sglang:cache_hit_rate The cache hit rate +# TYPE sglang:cache_hit_rate gauge +sglang:cache_hit_rate{model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.007507552643049313 +# HELP sglang:time_to_first_token_seconds Histogram of time to first token in seconds. +# TYPE sglang:time_to_first_token_seconds histogram +sglang:time_to_first_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 2.3518979474117756e+06 +sglang:time_to_first_token_seconds_bucket{le="0.001",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:time_to_first_token_seconds_bucket{le="0.005",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:time_to_first_token_seconds_bucket{le="0.01",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:time_to_first_token_seconds_bucket{le="0.02",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:time_to_first_token_seconds_bucket{le="0.04",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1.0 +sglang:time_to_first_token_seconds_bucket{le="0.06",model_name="meta-llama/Llama-3.1-8B-Instruct"} 3.0 +sglang:time_to_first_token_seconds_bucket{le="0.08",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="0.1",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="0.25",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="0.75",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="1.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 27.0 +sglang:time_to_first_token_seconds_bucket{le="2.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 140.0 +sglang:time_to_first_token_seconds_bucket{le="5.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 314.0 +sglang:time_to_first_token_seconds_bucket{le="7.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 941.0 +sglang:time_to_first_token_seconds_bucket{le="10.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1330.0 +sglang:time_to_first_token_seconds_bucket{le="15.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1970.0 +sglang:time_to_first_token_seconds_bucket{le="20.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 2326.0 +sglang:time_to_first_token_seconds_bucket{le="25.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 2417.0 +sglang:time_to_first_token_seconds_bucket{le="30.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 2513.0 +sglang:time_to_first_token_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11008.0 +sglang:time_to_first_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 11008.0 +# HELP sglang:e2e_request_latency_seconds Histogram of End-to-end request latency in seconds +# TYPE sglang:e2e_request_latency_seconds histogram +sglang:e2e_request_latency_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 3.116093850019932e+06 +sglang:e2e_request_latency_seconds_bucket{le="0.3",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:e2e_request_latency_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="0.8",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="1.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="1.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="2.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="2.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="5.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.0 +sglang:e2e_request_latency_seconds_bucket{le="10.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 10.0 +sglang:e2e_request_latency_seconds_bucket{le="15.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11.0 +sglang:e2e_request_latency_seconds_bucket{le="20.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 14.0 +sglang:e2e_request_latency_seconds_bucket{le="30.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 247.0 +sglang:e2e_request_latency_seconds_bucket{le="40.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 486.0 +sglang:e2e_request_latency_seconds_bucket{le="50.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 845.0 +sglang:e2e_request_latency_seconds_bucket{le="60.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1513.0 +sglang:e2e_request_latency_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11228.0 +sglang:e2e_request_latency_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 11228.0 +# HELP sglang:time_per_output_token_seconds Histogram of time per output token in seconds. +# TYPE sglang:time_per_output_token_seconds histogram +sglang:time_per_output_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 866964.5791549598 +sglang:time_per_output_token_seconds_bucket{le="0.005",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1.0 +sglang:time_per_output_token_seconds_bucket{le="0.01",model_name="meta-llama/Llama-3.1-8B-Instruct"} 73.0 +sglang:time_per_output_token_seconds_bucket{le="0.015",model_name="meta-llama/Llama-3.1-8B-Instruct"} 382.0 +sglang:time_per_output_token_seconds_bucket{le="0.02",model_name="meta-llama/Llama-3.1-8B-Instruct"} 593.0 +sglang:time_per_output_token_seconds_bucket{le="0.025",model_name="meta-llama/Llama-3.1-8B-Instruct"} 855.0 +sglang:time_per_output_token_seconds_bucket{le="0.03",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1035.0 +sglang:time_per_output_token_seconds_bucket{le="0.04",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1815.0 +sglang:time_per_output_token_seconds_bucket{le="0.05",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11685.0 +sglang:time_per_output_token_seconds_bucket{le="0.075",model_name="meta-llama/Llama-3.1-8B-Instruct"} 433413.0 +sglang:time_per_output_token_seconds_bucket{le="0.1",model_name="meta-llama/Llama-3.1-8B-Instruct"} 4.950195e+06 +sglang:time_per_output_token_seconds_bucket{le="0.15",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.039435e+06 +sglang:time_per_output_token_seconds_bucket{le="0.2",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.171662e+06 +sglang:time_per_output_token_seconds_bucket{le="0.3",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.266055e+06 +sglang:time_per_output_token_seconds_bucket{le="0.4",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.296752e+06 +sglang:time_per_output_token_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.312226e+06 +sglang:time_per_output_token_seconds_bucket{le="0.75",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.339675e+06 +sglang:time_per_output_token_seconds_bucket{le="1.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.357747e+06 +sglang:time_per_output_token_seconds_bucket{le="2.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.389414e+06 +sglang:time_per_output_token_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.400757e+06 +sglang:time_per_output_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.400757e+06 +# HELP sglang:func_latency_seconds Function latency in seconds +# TYPE sglang:func_latency_seconds histogram +sglang:func_latency_seconds_sum{name="generate_request"} 4.514771912145079 +sglang:func_latency_seconds_bucket{le="0.05",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.07500000000000001",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.1125",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.16875",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.253125",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.3796875",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.56953125",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.8542968750000001",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="1.2814453125",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="1.9221679687500002",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="2.8832519531250003",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="4.3248779296875",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="6.487316894531251",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="9.730975341796876",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="14.596463012695313",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="21.89469451904297",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="32.84204177856446",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="49.26306266784668",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="+Inf",name="generate_request"} 14007.0 +sglang:func_latency_seconds_count{name="generate_request"} 14007.0 +# HELP sglang:num_running_reqs The number of running requests +# TYPE sglang:num_running_reqs gauge +sglang:num_running_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct"} 162.0 +# HELP sglang:num_used_tokens The number of used tokens +# TYPE sglang:num_used_tokens gauge +sglang:num_used_tokens{model_name="meta-llama/Llama-3.1-8B-Instruct"} 123859.0 +# HELP sglang:gen_throughput The generate throughput (token/s) +# TYPE sglang:gen_throughput gauge +sglang:gen_throughput{model_name="meta-llama/Llama-3.1-8B-Instruct"} 86.50814177726902 +# HELP sglang:num_queue_reqs The number of requests in the waiting queue +# TYPE sglang:num_queue_reqs gauge +sglang:num_queue_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct"} 2826.0 +``` + +## Setup Guide + +This section describes how to set up the monitoring stack (Prometheus + Grafana) provided in the `examples/monitoring` directory. + +### Prerequisites + +- Docker and Docker Compose installed +- SGLang server running with metrics enabled + +### Usage + +1. **Start your SGLang server with metrics enabled:** + + ```bash + python -m sglang.launch_server --model-path --port 30000 --enable-metrics + ``` + Replace `` with the actual path to your model (e.g., `meta-llama/Meta-Llama-3.1-8B-Instruct`). Ensure the server is accessible from the monitoring stack (you might need `--host 0.0.0.0` if running in Docker). By default, the metrics endpoint will be available at `http://:30000/metrics`. + +2. **Navigate to the monitoring example directory:** + ```bash + cd examples/monitoring + ``` + +3. **Start the monitoring stack:** + ```bash + docker compose up -d + ``` + This command will start Prometheus and Grafana in the background. + +4. **Access the monitoring interfaces:** + * **Grafana:** Open your web browser and go to [http://localhost:3000](http://localhost:3000). + * **Prometheus:** Open your web browser and go to [http://localhost:9090](http://localhost:9090). + +5. **Log in to Grafana:** + * Default Username: `admin` + * Default Password: `admin` + You will be prompted to change the password upon your first login. + +6. **View the Dashboard:** + The SGLang dashboard is pre-configured and should be available automatically. Navigate to `Dashboards` -> `Browse` -> `SGLang Monitoring` folder -> `SGLang Dashboard`. + +### Troubleshooting + +* **Port Conflicts:** If you encounter errors like "port is already allocated," check if other services (including previous instances of Prometheus/Grafana) are using ports `9090` or `3000`. Use `docker ps` to find running containers and `docker stop ` to stop them, or use `lsof -i :` to find other processes using the ports. You might need to adjust the ports in the `docker-compose.yaml` file if they permanently conflict with other essential services on your system. + +To modify Grafana's port to the other one(like 3090) in your Docker Compose file, you need to explicitly specify the port mapping under the grafana service. + + Option 1: Add GF_SERVER_HTTP_PORT to the environment section: + ``` + environment: + - GF_AUTH_ANONYMOUS_ENABLED=true + - GF_SERVER_HTTP_PORT=3090 # <-- Add this line + ``` + Option 2: Use port mapping: + ``` + grafana: + image: grafana/grafana:latest + container_name: grafana + ports: + - "3090:3000" # <-- Host:Container port mapping + ``` +* **Connection Issues:** + * Ensure both Prometheus and Grafana containers are running (`docker ps`). + * Verify the Prometheus data source configuration in Grafana (usually auto-configured via `grafana/datasources/datasource.yaml`). Go to `Connections` -> `Data sources` -> `Prometheus`. The URL should point to the Prometheus service (e.g., `http://prometheus:9090`). + * Confirm that your SGLang server is running and the metrics endpoint (`http://:30000/metrics`) is accessible *from the Prometheus container*. If SGLang is running on your host machine and Prometheus is in Docker, use `host.docker.internal` (on Docker Desktop) or your machine's network IP instead of `localhost` in the `prometheus.yaml` scrape configuration. +* **No Data on Dashboard:** + * Generate some traffic to your SGLang server to produce metrics. For example, run a benchmark: + ```bash + python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 100 --random-input 128 --random-output 128 + ``` + * Check the Prometheus UI (`http://localhost:9090`) under `Status` -> `Targets` to see if the SGLang endpoint is being scraped successfully. + * Verify the `model_name` and `instance` labels in your Prometheus metrics match the variables used in the Grafana dashboard. You might need to adjust the Grafana dashboard variables or the labels in your Prometheus configuration. + +### Configuration Files + +The monitoring setup is defined by the following files within the `examples/monitoring` directory: + +* `docker-compose.yaml`: Defines the Prometheus and Grafana services. +* `prometheus.yaml`: Prometheus configuration, including scrape targets. +* `grafana/datasources/datasource.yaml`: Configures the Prometheus data source for Grafana. +* `grafana/dashboards/config/dashboard.yaml`: Tells Grafana to load dashboards from the specified path. +* `grafana/dashboards/json/sglang-dashboard.json`: The actual Grafana dashboard definition in JSON format. + +You can customize the setup by modifying these files. For instance, you might need to update the `static_configs` target in `prometheus.yaml` if your SGLang server runs on a different host or port. + +#### Check if the metrics are being collected + +Run `python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5` to generate some requests. + +Then you should be able to see the metrics in the Grafana dashboard. diff --git a/docs/references/torch_compile_cache.md b/docs/references/torch_compile_cache.md new file mode 100644 index 000000000..f2bb257f4 --- /dev/null +++ b/docs/references/torch_compile_cache.md @@ -0,0 +1,13 @@ +# Enabling cache for torch.compile + +SGLang uses `max-autotune-no-cudagraphs` mode of torch.compile. The auto-tuning can be slow. +If you want to deploy a model on many different machines, you can ship the torch.compile cache to these machines and skip the compilation steps. + +This is based on https://pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html + + +1. Generate the cache by setting TORCHINDUCTOR_CACHE_DIR and running the model once. +``` +TORCHINDUCTOR_CACHE_DIR=/root/inductor_root_cache python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile +``` +2. Copy the cache folder to other machines and launch the server with `TORCHINDUCTOR_CACHE_DIR`. diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 000000000..1a7e5d4eb --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,20 @@ +ipykernel +ipywidgets +jupyter_client +markdown>=3.4.0 +matplotlib +myst-parser +nbconvert +nbsphinx +pandoc +pillow +pydantic +sphinx +sphinx-book-theme +sphinx-copybutton +sphinx-tabs +nbstripout +sphinxcontrib-mermaid +urllib3<2.0.0 +gguf>=0.10.0 +sphinx-autobuild diff --git a/docs/serve.sh b/docs/serve.sh new file mode 100644 index 000000000..049f767cf --- /dev/null +++ b/docs/serve.sh @@ -0,0 +1,3 @@ +# Clean and serve documentation with auto-build +make clean +make serve diff --git a/docs/supported_models/embedding_models.md b/docs/supported_models/embedding_models.md new file mode 100644 index 000000000..437cb8284 --- /dev/null +++ b/docs/supported_models/embedding_models.md @@ -0,0 +1,87 @@ +# Embedding Models + +SGLang provides robust support for embedding models by integrating efficient serving mechanisms with its flexible programming interface. This integration allows for streamlined handling of embedding tasks, facilitating faster and more accurate retrieval and semantic search operations. SGLang's architecture enables better resource utilization and reduced latency in embedding model deployment. + +```{important} +Embedding models are executed with `--is-embedding` flag and some may require `--trust-remote-code` +``` + +## Quick Start + +### Launch Server + +```shell +python3 -m sglang.launch_server \ + --model-path Qwen/Qwen3-Embedding-4B \ + --is-embedding \ + --host 0.0.0.0 \ + --port 30000 +``` + +### Client Request + +```python +import requests + +url = "http://127.0.0.1:30000" + +payload = { + "model": "Qwen/Qwen3-Embedding-4B", + "input": "What is the capital of France?", + "encoding_format": "float" +} + +response = requests.post(url + "/v1/embeddings", json=payload).json() +print("Embedding:", response["data"][0]["embedding"]) +``` + + + +## Multimodal Embedding Example + +For multimodal models like GME that support both text and images: + +```shell +python3 -m sglang.launch_server \ + --model-path Alibaba-NLP/gme-Qwen2-VL-2B-Instruct \ + --is-embedding \ + --chat-template gme-qwen2-vl \ + --host 0.0.0.0 \ + --port 30000 +``` + +```python +import requests + +url = "http://127.0.0.1:30000" + +text_input = "Represent this image in embedding space." +image_path = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/023.jpg" + +payload = { + "model": "gme-qwen2-vl", + "input": [ + { + "text": text_input + }, + { + "image": image_path + } + ], +} + +response = requests.post(url + "/v1/embeddings", json=payload).json() + +print("Embeddings:", [x.get("embedding") for x in response.get("data", [])]) +``` + +## Supported Models + +| Model Family | Example Model | Chat Template | Description | +| ------------------------------------------ | -------------------------------------- | ------------- | --------------------------------------------------------------------------- | +| **E5 (Llama/Mistral based)** | `intfloat/e5-mistral-7b-instruct` | N/A | High-quality text embeddings based on Mistral/Llama architectures | +| **GTE-Qwen2** | `Alibaba-NLP/gte-Qwen2-7B-instruct` | N/A | Alibaba's text embedding model with multilingual support | +| **Qwen3-Embedding** | `Qwen/Qwen3-Embedding-4B` | N/A | Latest Qwen3-based text embedding model for semantic representation | +| **BGE** | `BAAI/bge-large-en-v1.5` | N/A | BAAI's text embeddings (requires `attention-backend` triton/torch_native) | +| **GME (Multimodal)** | `Alibaba-NLP/gme-Qwen2-VL-2B-Instruct`| `gme-qwen2-vl`| Multimodal embedding for text and image cross-modal tasks | +| **CLIP** | `openai/clip-vit-large-patch14-336` | N/A | OpenAI's CLIP for image and text embeddings | diff --git a/docs/supported_models/generative_models.md b/docs/supported_models/generative_models.md new file mode 100644 index 000000000..d6d3cdd45 --- /dev/null +++ b/docs/supported_models/generative_models.md @@ -0,0 +1,56 @@ +# Large Language Models + +These models accept text input and produce text output (e.g., chat completions). They are primarily large language models (LLMs), some with mixture-of-experts (MoE) architectures for scaling. + +## Example launch Command + +```shell +python3 -m sglang.launch_server \ + --model-path meta-llama/Llama-3.2-1B-Instruct \ # example HF/local path + --host 0.0.0.0 \ + --port 30000 \ +``` + +## Supported models + +Below the supported models are summarized in a table. + +If you are unsure if a specific architecture is implemented, you can search for it via GitHub. For example, to search for `Qwen3ForCausalLM`, use the expression: + +``` +repo:sgl-project/sglang path:/^python\/sglang\/srt\/models\// Qwen3ForCausalLM +``` + +in the GitHub search bar. + +| Model Family (Variants) | Example HuggingFace Identifier | Description | +|-------------------------------------|--------------------------------------------------|----------------------------------------------------------------------------------------| +| **DeepSeek** (v1, v2, v3/R1) | `deepseek-ai/DeepSeek-R1` | Series of advanced reasoning-optimized models (including a 671B MoE) trained with reinforcement learning; top performance on complex reasoning, math, and code tasks. [SGLang provides Deepseek v3/R1 model-specific optimizations](../basic_usage/deepseek.md) and [Reasoning Parser](../advanced_features/separate_reasoning.ipynb)| +| **GPT-OSS** | `openai/gpt-oss-20b`, `openai/gpt-oss-120b` | OpenAI’s latest GPT-OSS series for complex reasoning, agentic tasks, and versatile developer use cases.| +| **Qwen** (3, 3MoE, 3Next, 2.5, 2 series) | `Qwen/Qwen3-0.6B`, `Qwen/Qwen3-30B-A3B` `Qwen/Qwen3-Next-80B-A3B-Instruct ` | Alibaba’s latest Qwen3 series for complex reasoning, language understanding, and generation tasks; Support for MoE variants along with previous generation 2.5, 2, etc. [SGLang provides Qwen3 specific reasoning parser](../advanced_features/separate_reasoning.ipynb)| +| **Llama** (2, 3.x, 4 series) | `meta-llama/Llama-4-Scout-17B-16E-Instruct` | Meta's open LLM series, spanning 7B to 400B parameters (Llama 2, 3, and new Llama 4) with well-recognized performance. [SGLang provides Llama-4 model-specific optimizations](../basic_usage/llama4.md) | +| **Mistral** (Mixtral, NeMo, Small3) | `mistralai/Mistral-7B-Instruct-v0.2` | Open 7B LLM by Mistral AI with strong performance; extended into MoE (“Mixtral”) and NeMo Megatron variants for larger scale. | +| **Gemma** (v1, v2, v3) | `google/gemma-3-1b-it` | Google’s family of efficient multilingual models (1B–27B); Gemma 3 offers a 128K context window, and its larger (4B+) variants support vision input. | +| **Phi** (Phi-1.5, Phi-2, Phi-3, Phi-4, Phi-MoE series) | `microsoft/Phi-4-multimodal-instruct`, `microsoft/Phi-3.5-MoE-instruct` | Microsoft’s Phi family of small models (1.3B–5.6B); Phi-4-multimodal (5.6B) processes text, images, and speech, Phi-4-mini is a high-accuracy text model and Phi-3.5-MoE is a mixture-of-experts model. | +| **MiniCPM** (v3, 4B) | `openbmb/MiniCPM3-4B` | OpenBMB’s series of compact LLMs for edge devices; MiniCPM 3 (4B) achieves GPT-3.5-level results in text tasks. | +| **OLMoE** (Open MoE) | `allenai/OLMoE-1B-7B-0924` | Allen AI’s open Mixture-of-Experts model (7B total, 1B active parameters) delivering state-of-the-art results with sparse expert activation. | +| **StableLM** (3B, 7B) | `stabilityai/stablelm-tuned-alpha-7b` | StabilityAI’s early open-source LLM (3B & 7B) for general text generation; a demonstration model with basic instruction-following ability. | +| **Command-R** (Cohere) | `CohereForAI/c4ai-command-r-v01` | Cohere’s open conversational LLM (Command series) optimized for long context, retrieval-augmented generation, and tool use. | +| **DBRX** (Databricks) | `databricks/dbrx-instruct` | Databricks’ 132B-parameter MoE model (36B active) trained on 12T tokens; competes with GPT-3.5 quality as a fully open foundation model. | +| **Grok** (xAI) | `xai-org/grok-1` | xAI’s grok-1 model known for vast size(314B parameters) and high quality; integrated in SGLang for high-performance inference. | +| **ChatGLM** (GLM-130B family) | `THUDM/chatglm2-6b` | Zhipu AI’s bilingual chat model (6B) excelling at Chinese-English dialogue; fine-tuned for conversational quality and alignment. | +| **InternLM 2** (7B, 20B) | `internlm/internlm2-7b` | Next-gen InternLM (7B and 20B) from SenseTime, offering strong reasoning and ultra-long context support (up to 200K tokens). | +| **ExaONE 3** (Korean-English) | `LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct` | LG AI Research’s Korean-English model (7.8B) trained on 8T tokens; provides high-quality bilingual understanding and generation. | +| **Baichuan 2** (7B, 13B) | `baichuan-inc/Baichuan2-13B-Chat` | BaichuanAI’s second-generation Chinese-English LLM (7B/13B) with improved performance and an open commercial license. | +| **XVERSE** (MoE) | `xverse/XVERSE-MoE-A36B` | Yuanxiang’s open MoE LLM (XVERSE-MoE-A36B: 255B total, 36B active) supporting ~40 languages; delivers 100B+ dense-level performance via expert routing. | +| **SmolLM** (135M–1.7B) | `HuggingFaceTB/SmolLM-1.7B` | Hugging Face’s ultra-small LLM series (135M–1.7B params) offering surprisingly strong results, enabling advanced AI on mobile/edge devices. | +| **GLM-4** (Multilingual 9B) | `ZhipuAI/glm-4-9b-chat` | Zhipu’s GLM-4 series (up to 9B parameters) – open multilingual models with support for 1M-token context and even a 5.6B multimodal variant (Phi-4V). | +| **MiMo** (7B series) | `XiaomiMiMo/MiMo-7B-RL` | Xiaomi's reasoning-optimized model series, leverages Multiple-Token Prediction for faster inference. | +| **ERNIE-4.5** (4.5, 4.5MoE series) | `baidu/ERNIE-4.5-21B-A3B-PT` | Baidu's ERNIE-4.5 series which consists of MoE with 47B and 3B active parameters, with the largest model having 424B total parameters, as well as a 0.3B dense model. | +| **Arcee AFM-4.5B** | `arcee-ai/AFM-4.5B-Base` | Arcee's foundational model series for real world reliability and edge deployments. | +| **Persimmon** (8B) | `adept/persimmon-8b-chat` | Adept’s open 8B model with a 16K context window and fast inference; trained for broad usability and licensed under Apache 2.0. | +| **Ling** (16.8B–290B) | `inclusionAI/Ling-lite`, `inclusionAI/Ling-plus` | InclusionAI’s open MoE models. Ling-Lite has 16.8B total / 2.75B active parameters, and Ling-Plus has 290B total / 28.8B active parameters. They are designed for high performance on NLP and complex reasoning tasks. | +| **Granite 3.0, 3.1** (IBM) | `ibm-granite/granite-3.1-8b-instruct` | IBM's open dense foundation models optimized for reasoning, code, and business AI use cases. Integrated with Red Hat and watsonx systems. | +| **Granite 3.0 MoE** (IBM) | `ibm-granite/granite-3.0-3b-a800m-instruct` | IBM’s Mixture-of-Experts models offering strong performance with cost-efficiency. MoE expert routing designed for enterprise deployment at scale. | +| **Llama Nemotron Super** (v1, v1.5, NVIDIA) | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, `nvidia/Llama-3_3-Nemotron-Super-49B-v1_5` | The [NVIDIA Nemotron](https://www.nvidia.com/en-us/ai-data-science/foundation-models/nemotron/) family builds on the strongest open models in the ecosystem by enhancing them with greater accuracy, efficiency, and transparency using NVIDIA open synthetic datasets, advanced techniques, and tools. This enables the creation of practical, right-sized, and high-performing AI agents. | +| **Llama Nemotron Ultra** (v1, NVIDIA) | `nvidia/Llama-3_1-Nemotron-Ultra-253B-v1` | The [NVIDIA Nemotron](https://www.nvidia.com/en-us/ai-data-science/foundation-models/nemotron/) family builds on the strongest open models in the ecosystem by enhancing them with greater accuracy, efficiency, and transparency using NVIDIA open synthetic datasets, advanced techniques, and tools. This enables the creation of practical, right-sized, and high-performing AI agents. | diff --git a/docs/supported_models/modelscope.md b/docs/supported_models/modelscope.md new file mode 100644 index 000000000..4740c2770 --- /dev/null +++ b/docs/supported_models/modelscope.md @@ -0,0 +1,28 @@ +# Use Models From ModelScope + +To use a model from [ModelScope](https://www.modelscope.cn), set the environment variable `SGLANG_USE_MODELSCOPE`. + +```bash +export SGLANG_USE_MODELSCOPE=true +``` + +We take [Qwen2-7B-Instruct](https://www.modelscope.cn/models/qwen/qwen2-7b-instruct) as an example. + +Launch the Server: +```bash +python -m sglang.launch_server --model-path qwen/Qwen2-7B-Instruct --port 30000 +``` + +Or start it by docker: + +```bash +docker run --gpus all \ + -p 30000:30000 \ + -v ~/.cache/modelscope:/root/.cache/modelscope \ + --env "SGLANG_USE_MODELSCOPE=true" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct --host 0.0.0.0 --port 30000 +``` + +Note that modelscope uses a different cache directory than huggingface. You may need to set it manually to avoid running out of disk space. diff --git a/docs/supported_models/multimodal_language_models.md b/docs/supported_models/multimodal_language_models.md new file mode 100644 index 000000000..a2adf99cb --- /dev/null +++ b/docs/supported_models/multimodal_language_models.md @@ -0,0 +1,42 @@ +# Multimodal Language Models + +These models accept multi-modal inputs (e.g., images and text) and generate text output. They augment language models with multimodal encoders. + +## Example launch Command + +```shell +python3 -m sglang.launch_server \ + --model-path meta-llama/Llama-3.2-11B-Vision-Instruct \ # example HF/local path + --host 0.0.0.0 \ + --port 30000 \ +``` + +## Supported models + +Below the supported models are summarized in a table. + +If you are unsure if a specific architecture is implemented, you can search for it via GitHub. For example, to search for `Qwen2_5_VLForConditionalGeneration`, use the expression: + +``` +repo:sgl-project/sglang path:/^python\/sglang\/srt\/models\// Qwen2_5_VLForConditionalGeneration +``` + +in the GitHub search bar. + + +| Model Family (Variants) | Example HuggingFace Identifier | Chat Template | Description | +|----------------------------|--------------------------------------------|------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| **Qwen-VL** (Qwen2 series) | `Qwen/Qwen2.5-VL-7B-Instruct` | `qwen2-vl` | Alibaba’s vision-language extension of Qwen; for example, Qwen2.5-VL (7B and larger variants) can analyze and converse about image content. | +| **DeepSeek-VL2** | `deepseek-ai/deepseek-vl2` | `deepseek-vl2` | Vision-language variant of DeepSeek (with a dedicated image processor), enabling advanced multimodal reasoning on image and text inputs. | +| **Janus-Pro** (1B, 7B) | `deepseek-ai/Janus-Pro-7B` | `janus-pro` | DeepSeek’s open-source multimodal model capable of both image understanding and generation. Janus-Pro employs a decoupled architecture for separate visual encoding paths, enhancing performance in both tasks. | +| **MiniCPM-V / MiniCPM-o** | `openbmb/MiniCPM-V-2_6` | `minicpmv` | MiniCPM-V (2.6, ~8B) supports image inputs, and MiniCPM-o adds audio/video; these multimodal LLMs are optimized for end-side deployment on mobile/edge devices. | +| **Llama 3.2 Vision** (11B) | `meta-llama/Llama-3.2-11B-Vision-Instruct` | `llama_3_vision` | Vision-enabled variant of Llama 3 (11B) that accepts image inputs for visual question answering and other multimodal tasks. | +| **LLaVA** (v1.5 & v1.6) | *e.g.* `liuhaotian/llava-v1.5-13b` | `vicuna_v1.1` | Open vision-chat models that add an image encoder to LLaMA/Vicuna (e.g. LLaMA2 13B) for following multimodal instruction prompts. | +| **LLaVA-NeXT** (8B, 72B) | `lmms-lab/llava-next-72b` | `chatml-llava` | Improved LLaVA models (with an 8B Llama3 version and a 72B version) offering enhanced visual instruction-following and accuracy on multimodal benchmarks. | +| **LLaVA-OneVision** | `lmms-lab/llava-onevision-qwen2-7b-ov` | `chatml-llava` | Enhanced LLaVA variant integrating Qwen as the backbone; supports multiple images (and even video frames) as inputs via an OpenAI Vision API-compatible format. | +| **Gemma 3 (Multimodal)** | `google/gemma-3-4b-it` | `gemma-it` | Gemma 3's larger models (4B, 12B, 27B) accept images (each image encoded as 256 tokens) alongside text in a combined 128K-token context. | +| **Kimi-VL** (A3B) | `moonshotai/Kimi-VL-A3B-Instruct` | `kimi-vl` | Kimi-VL is a multimodal model that can understand and generate text from images. | +| **Mistral-Small-3.1-24B** | `mistralai/Mistral-Small-3.1-24B-Instruct-2503` | `mistral` | Mistral 3.1 is a multimodal model that can generate text from text or images input. It also supports tool calling and structured output. | +| **Phi-4-multimodal-instruct** | `microsoft/Phi-4-multimodal-instruct` | `phi-4-mm` | Phi-4-multimodal-instruct is the multimodal variant of the Phi-4-mini model, enhanced with LoRA for improved multimodal capabilities. It supports text, vision and audio modalities in SGLang. | +| **MiMo-VL** (7B) | `XiaomiMiMo/MiMo-VL-7B-RL` | `mimo-vl` | Xiaomi's compact yet powerful vision-language model featuring a native resolution ViT encoder for fine-grained visual details, an MLP projector for cross-modal alignment, and the MiMo-7B language model optimized for complex reasoning tasks. | +| **GLM-4.5V** (106B) / **GLM-4.1V**(9B) | `zai-org/GLM-4.5V` | `glm-4v` | GLM-4.5V and GLM-4.1V-Thinking: Towards Versatile Multimodal Reasoning with Scalable Reinforcement Learning | diff --git a/docs/supported_models/rerank_models.md b/docs/supported_models/rerank_models.md new file mode 100644 index 000000000..b6f2ffa20 --- /dev/null +++ b/docs/supported_models/rerank_models.md @@ -0,0 +1,49 @@ +# Rerank Models + +SGLang offers comprehensive support for rerank models by incorporating optimized serving frameworks with a flexible programming interface. This setup enables efficient processing of cross-encoder reranking tasks, improving the accuracy and relevance of search result ordering. SGLang’s design ensures high throughput and low latency during reranker model deployment, making it ideal for semantic-based result refinement in large-scale retrieval systems. + +```{important} +They are executed with `--is-embedding` and some may require `--trust-remote-code` +``` + +## Example Launch Command + +```shell +python3 -m sglang.launch_server \ + --model-path BAAI/bge-reranker-v2-m3 \ + --host 0.0.0.0 \ + --disable-radix-cache \ + --chunked-prefill-size -1 \ + --attention-backend triton \ + --is-embedding \ + --port 30000 +``` + +## Example Client Request + +```python +import requests + +url = "http://127.0.0.1:30000/v1/rerank" + +payload = { + "model": "BAAI/bge-reranker-v2-m3", + "query": "what is panda?", + "documents": [ + "hi", + "The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China." + ] +} + +response = requests.post(url, json=payload) +response_json = response.json() + +for item in response_json: + print(f"Score: {item['score']:.2f} - Document: '{item['document']}'") +``` + +## Supported rerank models + +| Model Family (Rerank) | Example HuggingFace Identifier | Chat Template | Description | +|------------------------------------------------|--------------------------------------|---------------|----------------------------------------------------------------------------------------------------------------------------------| +| **BGE-Reranker (BgeRerankModel)** | `BAAI/bge-reranker-v2-m3` | N/A | Currently only support `attention-backend` `triton` and `torch_native`. high-performance cross-encoder reranker model from BAAI. Suitable for reranking search results based on semantic relevance. | diff --git a/docs/supported_models/reward_models.md b/docs/supported_models/reward_models.md new file mode 100644 index 000000000..ef4474637 --- /dev/null +++ b/docs/supported_models/reward_models.md @@ -0,0 +1,28 @@ +# Reward Models + +These models output a scalar reward score or classification result, often used in reinforcement learning or content moderation tasks. + +```{important} +They are executed with `--is-embedding` and some may require `--trust-remote-code`. +``` + +## Example launch Command + +```shell +python3 -m sglang.launch_server \ + --model-path Qwen/Qwen2.5-Math-RM-72B \ # example HF/local path + --is-embedding \ + --host 0.0.0.0 \ + --tp-size=4 \ # set for tensor parallelism + --port 30000 \ +``` + +## Supported models + +| Model Family (Reward) | Example HuggingFace Identifier | Description | +|---------------------------------------------------------------------------|-----------------------------------------------------|---------------------------------------------------------------------------------| +| **Llama (3.1 Reward / `LlamaForSequenceClassification`)** | `Skywork/Skywork-Reward-Llama-3.1-8B-v0.2` | Reward model (preference classifier) based on Llama 3.1 (8B) for scoring and ranking responses for RLHF. | +| **Gemma 2 (27B Reward / `Gemma2ForSequenceClassification`)** | `Skywork/Skywork-Reward-Gemma-2-27B-v0.2` | Derived from Gemma‑2 (27B), this model provides human preference scoring for RLHF and multilingual tasks. | +| **InternLM 2 (Reward / `InternLM2ForRewardMode`)** | `internlm/internlm2-7b-reward` | InternLM 2 (7B)–based reward model used in alignment pipelines to guide outputs toward preferred behavior. | +| **Qwen2.5 (Reward - Math / `Qwen2ForRewardModel`)** | `Qwen/Qwen2.5-Math-RM-72B` | A 72B math-specialized RLHF reward model from the Qwen2.5 series, tuned for evaluating and refining responses. | +| **Qwen2.5 (Reward - Sequence / `Qwen2ForSequenceClassification`)** | `jason9693/Qwen2.5-1.5B-apeach` | A smaller Qwen2.5 variant used for sequence classification, offering an alternative RLHF scoring mechanism. | diff --git a/docs/supported_models/support_new_models.md b/docs/supported_models/support_new_models.md new file mode 100644 index 000000000..511a8f398 --- /dev/null +++ b/docs/supported_models/support_new_models.md @@ -0,0 +1,320 @@ +# How to Support New Models + +This document explains how to add support for new language models and multimodal large language models (MLLMs) in +SGLang. It also covers how to test new models and register external implementations. + +## How to Support a New Language Model + +To support a new model in SGLang, you only need to add a single file under +the [SGLang Models Directory](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/models). You can learn +from existing model implementations and create a new file for your model. For most models, you should be able to find a +similar model to start with (e.g., starting from Llama). Also refer how +to [port a Model from vLLM to SGLang](#port-a-model-from-vllm-to-sglang) + +## How to Support a New Multimodal Large Language Model + +To support a new multimodal large language model (MLLM) in SGLang, there are several key components in addition to the +standard LLM support: + +1. **Register your new model as multimodal**: + Extend `is_multimodal_model` + in [model_config.py](https://github.com/sgl-project/sglang/blob/0ab3f437aba729b348a683ab32b35b214456efc7/python/sglang/srt/configs/model_config.py#L561) + to return `True` for your model. + +2. **Register a new chat-template**: + Only when your default chat-template is unable to accept images as input: Register a new chat template in [conversation.py](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/conversation.py) and the corresponding matching function. + +3. **Multimodal Data Processor**: + Define a new `Processor` class that inherits from `BaseMultimodalProcessor` and register this processor as your + model’s dedicated processor. + See [multimodal_processor.py](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/multimodal/processors) + for more details. + +4. **Handle Multimodal Tokens**: + Implement a `pad_input_ids` function for your new model. In this function, multimodal tokens in the prompt should be + expanded (if necessary) and padded with multimodal-data-hashes so that SGLang can recognize different multimodal data + with `RadixAttention`. + +5. **Handle Image Feature Extraction**: + Implement a `get_image_feature` function for your new model, which extracts image features from raw image data and converts them into the embeddings used by the language model. + +6. **Adapt to Vision Attention**: + Adapt the multi-headed `Attention` of ViT with SGLang’s `VisionAttention`. + +You can refer to [Qwen2VL](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/qwen2_vl.py) or +other mllm implementations. These models demonstrate how to correctly handle both multimodal and textual inputs. + +## Testing and Debugging + +Please note all your testing and benchmarking results in PR description. + +### Interactive Debugging + +For interactive debugging, compare the outputs of Hugging Face/Transformers and SGLang. The following two commands +should give the same text output and very similar prefill logits: + +- Get the reference output: + ```bash + python3 scripts/playground/reference_hf.py --model-path [new model] --model-type {text,mllm} + ``` +- Get the SGLang output: + ```bash + python3 -m sglang.bench_one_batch --correct --model [new model] + ``` + +### Add the Model to the Test Suite + +To ensure the new model is well maintained, add it to the test suite by including it in the `ALL_OTHER_MODELS` list in +the [test_generation_models.py](https://github.com/sgl-project/sglang/blob/main/test/srt/models/test_generation_models.py) +file, test the new model on your local machine and report the results on demonstrative benchmarks (GSM8K, MMLU, MMMU, +MMMU-Pro, etc.) in your PR. \\ +For VLMs, also include a test in `test_vision_openai_server_{x}.py` (e.g. [test_vision_openai_server_a.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_vision_openai_server_a.py), [test_vision_openai_server_b.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_vision_openai_server_b.py)). + + +This is an example command to run to test a new model on your local machine: + +```bash +ONLY_RUN=Qwen/Qwen2-1.5B python3 -m unittest test_generation_models.TestGenerationModels.test_others +``` + +### Benchmark + +- **(Required) MMMU**: follow MMMU benchmark [README.md](https://github.com/sgl-project/sglang/blob/main/benchmark/mmmu/README.md) to get SGLang vs. HF Transformer accuracy comparison. The accuracy score from SGLang run should not be much lower than that from HF Transformer run. Similarly, follow https://docs.sglang.ai/developer_guide/benchmark_and_profiling.html to get performance comparison: TTFT and throughput must meet or exceed baselines (e.g., HF Transformer). +- **(Optional) Other evals**: If you ran other evals, please note the results in PR description. + +## Port a Model from vLLM to SGLang + +The [vLLM Models Directory](https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models) is a valuable +resource, as vLLM covers many models. SGLang reuses vLLM’s interface and some layers, making it easier to port models +from vLLM to SGLang. + +To port a model from vLLM to SGLang: + +- Compare these two files for guidance: + - [SGLang Llama Implementation](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/llama.py) + - [vLLM Llama Implementation](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama.py) +- The major differences include: + - **Replace vLLM’s `Attention` with `RadixAttention`** (ensure you pass `layer_id` to `RadixAttention`). + - **Replace vLLM’s `LogitsProcessor` with SGLang’s `LogitsProcessor`.** + - **Replace the multi-headed `Attention` of ViT with SGLang’s `VisionAttention`.** + - **Replace other vLLM layers** (such as `RMSNorm`, `SiluAndMul`) with SGLang layers. + - **Remove `Sample`.** + - **Change the `forward()` functions** and add a `forward_batch()` method. + - **Add `EntryClass`** at the end. + - **Ensure that the new implementation uses only SGLang components** and does not rely on any vLLM components. + +Note: make sure you add your new model to the supported models list in the supported models documentation. + +## Registering an External Model Implementation + +In addition to the methods above, you can register your new model with the `ModelRegistry` before launching the server. +This allows you to integrate your model without modifying the source code. + +For example: + +```python +from sglang.srt.models.registry import ModelRegistry +from sglang.srt.entrypoints.http_server import launch_server + +# For a single model, add it to the registry: +ModelRegistry.models[model_name] = model_class + +# For multiple models, you can imitate the import_model_classes() function: +from functools import lru_cache + +@lru_cache() +def import_new_model_classes(): + model_arch_name_to_cls = {} + # Populate model_arch_name_to_cls with your new model classes. + ... + return model_arch_name_to_cls + +ModelRegistry.models.update(import_new_model_classes()) + +# Launch the server with your server arguments: +launch_server(server_args) +``` + +## Example: Implementing and Serving a Llama Wrapper Model + +Below is an introductory, step-by-step walkthrough on how to implement a new model end-to-end in SGLang and then run it via the [Offline Engine](https://github.com/sgl-project/sglang/blob/main/docs/basic_usage/offline_engine_api.ipynb). + +### Implementing Our Model + +To keep things simple, this new model will be a simple wrapper around [Llama 3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct), and our goal will be just to bias the output logits for each `forward` call by taking the square root of each individual logit. + +Let's start by defining our model in a file called `llama_wrapper.py`. +The first step is to import the necessary libraries from SRT, which is SGLang's internal backend. + +```python +# In the file `llama_wrapper.py` + +import torch +from transformers import LlamaConfig +from typing import Optional +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors + +from sglang.srt.models.llama import LlamaForCausalLM +``` + +Next, we declare a new `class` for our model and have it inherit from `LlamaForCausalLM`, which allows our model to access `LlamaForCausalLM`'s predefined modules and layers, such as `LlamaAttention` and `LlamaMLP`. +Note that almost all model implementations take in `config` and `quant_config` as arguments for their `__init__` method; `config` and `quant_config` are passed in via [`model_loader/loader.py`](https://github.com/sgl-project/sglang/blob/bf72b80122fd888bf619d17b96fa3e323ab809fc/python/sglang/srt/model_loader/loader.py#L219). +Because we have inherited from `LlamaForCausalLM`, we can pass our parameters directly to its constructor, which will set the member variables for us. + +```python +class LlamaWrapper(LlamaForCausalLM): + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config=config, quant_config=quant_config, prefix=prefix) +``` + +Now, we want to define the `forward` method, which is what will be called at inference time. +Note that the signature for `forward` is essentially the same for any model; you can take a look at the other models defined in the [`models` directory](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/) for references. +To see where exactly `forward` is called in the SGLang runtime's internals, take a look at [`forward_decode`](https://github.com/sgl-project/sglang/blob/bf72b80122fd888bf619d17b96fa3e323ab809fc/python/sglang/srt/model_executor/model_runner.py#L1705) and [`forward_extend`](https://github.com/sgl-project/sglang/blob/bf72b80122fd888bf619d17b96fa3e323ab809fc/python/sglang/srt/model_executor/model_runner.py#L1724) in the [`ModelRunner` class](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/model_executor/model_runner.py). + +```python + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + input_embeds: Optional[torch.Tensor] = None, + get_embedding: bool = False, + ) -> LogitsProcessorOutput: +``` + +We now call the `__call__` method for `self.model` (which is a member variable that `LlamaForCausalLM` defines in its `__init__` method), which eventually calls `LlamaForCausalLM`'s `forward` method. +After that, we feed the `hidden_states` into our model's `LogitsProcessor` (again defined in `LlamaForCausalLM`). + +```python + hidden_states = self.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) + + res: LogitsProcessorOutput = self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + forward_batch, + ) +``` + +After receiving the logits for the next token, we can finally perform our biasing step. + +```python + orig_logits = res.next_token_logits + res.next_token_logits = torch.where( + orig_logits > 0, + orig_logits.sqrt(), + orig_logits + ) + + return res +``` +Now, our `LlamaWrapper` model is created and ready to be served! + +### Serving Our Model Via SGLang's Offline Engine + +The next step of this walkthrough involves hosting our new model offline, so that it can be served locally and without an HTTP server. + +First, create a new file called `run.py`. +Now, we must ensure that SGLang's `ModelRegistry` can find our model. +To do this, we first download the model's configuration and weights from Huggingface. + +```python +# In the file `run.py` + +import asyncio +from functools import lru_cache +from huggingface_hub import snapshot_download +from llama_wrapper import LlamaWrapper # Make sure to import our new model! +import sglang as sgl +from sglang.srt.models.registry import ModelRegistry + +# Make sure to request access to this model on Huggingface, then export your +# `HF_TOKEN` to download the model snapshot +llama_dir = snapshot_download( + repo_id="meta-llama/Llama-3.1-8B-Instruct", + local_dir="./llama_ckpt", +) +``` + +Now that we have our model on disk, we want to point it to `LlamaWrapper` by changing the `architectures` field in `./llama_ckpt/config.json` to be `LlamaWrapper`. +That way, when we pass in the path of our model checkpoint to SGLang, it will know that we want to use "LlamaWrapper" instead of "LlamaForCausalLM" as our model. + +```python +{ + "architectures": [ + # "LlamaForCausalLM" + "LlamaWrapper" + ], + ... +} +``` + +However, if we don't link our `LlamaWrapper` class to the "LlamaWrapper" registry keyword, then SGLang won't be able to find our model. +Thus, to register our `LlamaWrapper`, we want to follow the steps in the above section titled "Registering an External Model Implementation". + +```python +@lru_cache() +def import_new_model_classes(): + model_arch_name_to_cls = {"LlamaWrapper": LlamaWrapper} + return model_arch_name_to_cls + +ModelRegistry.models.update(import_new_model_classes()) +``` + +Lastly, when we create our `Engine`, we just pass in the path to the local model directory. +Then, our `LlamaWrapper` is ready to be served; for this walkthrough, we will use SGLang `Engine`'s non-streaming asynchronous generation endpoint. + +```python +def main(): + llm = sgl.Engine(model_path="./llama_ckpt") + sampling_params = {"temperature": 0.2, "top_k": 5} + prompts = [ + "Write a short, neutral self-introduction for a fictional character. Hello, my name is", + "Provide a concise factual statement about France’s capital city. The capital of France is", + "Explain possible future trends in artificial intelligence. The future of AI is", + ] + + asyncio.run(run_llm(llm, sampling_params, prompts)) + + llm.shutdown() + +async def run_llm( + llm, + sampling_params, + prompts, +) -> None: + outputs = await llm.async_generate(prompts, sampling_params) + + for prompt, output in zip(prompts, outputs): + print(f"\nPrompt: {prompt}") + print(f"Generated text: {output['text']}") + +if __name__ == "__main__": + main() +``` + +Now, when we call `python run.py`, we will get the outputs of our newly created model! + + +## Documentation +Add to table of supported models in [generative_models.md](https://github.com/sgl-project/sglang/blob/main/docs/supported_models/generative_models.md) or [multimodal_language_models.md](https://github.com/sgl-project/sglang/blob/main/docs/supported_models/multimodal_language_models.md) + +--- + +By following these guidelines, you can add support for new language models and multimodal large language models in +SGLang and ensure they are thoroughly tested and easily integrated into the system. diff --git a/docs/supported_models/transformers_fallback.md b/docs/supported_models/transformers_fallback.md new file mode 100644 index 000000000..3c7dd961c --- /dev/null +++ b/docs/supported_models/transformers_fallback.md @@ -0,0 +1,58 @@ +# Transformers fallback in SGLang + +`sglang` can fall back to using models that are available in `transformers`. This works for most decoder-style language models and support for vision-language models is coming soon! + +## Example launch Command + +By default, we will use sglang implementation if it is available. Otherwise, we will fall back to transformers one. However, you can switch the implementation by setting `--model-impl` to `transformers`. + +```shell +python3 -m sglang.launch_server \ + --model-path meta-llama/Llama-3.2-1B-Instruct \ + --host 0.0.0.0 \ + --port 30000 \ + --model-impl transformers +``` + +## Supported features + +### Quantization + +Transformers fall back has supported most of available quantization in SGLang (except GGUF). See [Quantization page](../advanced_features/quantization.md) for more information about supported quantization in SGLang. + +### Remote code + +This fallback also means that any model on the hub that can be used in `transformers` with `trust_remote_code=True` that correctly implements attention can be used in production! + +A model just needs the following two things: + +```python +from transformers import PreTrainedModel +from torch import nn + +class MyAttention(nn.Module): + + def forward(self, hidden_states, **kwargs): # <- kwargs are required + + ... + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + **kwargs, + ) + ... + +class MyModel(PreTrainedModel): + _supports_attention_backend = True +``` + +Here is what happens in the background: + +1. The config is loaded +2. `MyModel` python class is loaded from the `auto_map`, and we check that the model `_supports_attention_backend`. +3. The `TransformersModel` backend is used. See `/srt/models/transformers`, which leverages `self.config._attn_implementation = "sglang"`, thus the need to use `ALL_ATTENTION_FUNCTIONS`. + +That's it! diff --git a/docs/wrap_run_llm.py b/docs/wrap_run_llm.py new file mode 100644 index 000000000..2d21442cf --- /dev/null +++ b/docs/wrap_run_llm.py @@ -0,0 +1,47 @@ +import os +import re + + +def insert_runllm_widget(html_content): + # RunLLM Widget script to be inserted + widget_script = """ + + + """ + + # Find the closing body tag and insert the widget script before it + return re.sub(r"", f"{widget_script}\n", html_content) + + +def process_html_files(build_dir): + for root, dirs, files in os.walk(build_dir): + for file in files: + if file.endswith(".html"): + file_path = os.path.join(root, file) + + # Read the HTML file + with open(file_path, "r", encoding="utf-8") as f: + content = f.read() + + # Insert the RunLLM widget + modified_content = insert_runllm_widget(content) + + # Write back the modified content + with open(file_path, "w", encoding="utf-8") as f: + f.write(modified_content) + + +def main(): + # Get the build directory path + build_dir = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "_build", "html" + ) + # Process all HTML files + if os.path.exists(build_dir): + process_html_files(build_dir) + else: + print(f"Build directory not found: {build_dir}") + + +if __name__ == "__main__": + main() diff --git a/examples/chat_template/tool_chat_template_deepseekr1.jinja b/examples/chat_template/tool_chat_template_deepseekr1.jinja new file mode 100644 index 000000000..bb86cbd9e --- /dev/null +++ b/examples/chat_template/tool_chat_template_deepseekr1.jinja @@ -0,0 +1,92 @@ +{% if not add_generation_prompt is defined %} + {% set add_generation_prompt = false %} +{% endif %} +{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true, is_last_user=false) %} +{%- for message in messages %} + {%- if message['role'] == 'system' %} + {%- if ns.is_first_sp %} + {% set ns.system_prompt = ns.system_prompt + message['content'] %} + {% set ns.is_first_sp = false %} + {%- else %} + {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %} + {%- endif %} + {%- endif %} +{%- endfor %} + +{# --- Append tool descriptions if tools are defined --- #} +{% if tools is defined and tools is not none %} + {% set tool_ns = namespace(text='You are a helpful assistant with tool calling capabilities. ' + 'When a tool call is needed, you MUST use the following format to issue the call:\n' + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>FUNCTION_NAME\n' + '```json\n{"param1": "value1", "param2": "value2"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>\n\n' + 'Make sure the JSON is valid.' + '## Tools\n\n### Function\n\nYou have the following functions available:\n\n') %} + {% for tool in tools %} + {% set tool_ns.text = tool_ns.text + '- `' + tool['name'] + '`:\n```json\n' + (tool | tojson) + '\n```\n' %} + {% endfor %} + {% set ns.system_prompt = ns.system_prompt + '\n\n' + tool_ns.text %} +{% endif %} + +{{ bos_token }} +{{ ns.system_prompt }} +{%- for message in messages %} + {% set content = message['content'] %} + {%- if message['role'] == 'user' %} + {%- set ns.is_tool = false -%} + {%- set ns.is_first = false -%} + {%- set ns.is_last_user = true -%} + {{'<｜User｜>' + content + '<｜Assistant｜>'}} + {%- endif %} + {%- if message['role'] == 'assistant' %} + {% if '' in content %} + {% set content = content.split('')[-1] %} + {% endif %} + {% endif %} + {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %} + {%- set ns.is_last_user = false -%} + {%- if ns.is_tool %} + {{'<｜tool▁outputs▁end｜>'}} + {%- endif %} + {%- set ns.is_first = false %} + {%- set ns.is_tool = false -%} + {%- set ns.is_output_first = true %} + {%- for tool in message['tool_calls'] %} + {%- if not ns.is_first %} + {%- if content is none %} + {{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}} + {%- else %} + {{content + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}} + {%- endif %} + {%- set ns.is_first = true -%} + {%- else %} + {{'\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}} + {%- endif %} + {%- endfor %} + {{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}} + {%- endif %} + {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none)%} + {%- set ns.is_last_user = false -%} + {%- if ns.is_tool %} + {{'<｜tool▁outputs▁end｜>' + content + '<｜end▁of▁sentence｜>'}} + {%- set ns.is_tool = false -%} + {%- else %} + {{content + '<｜end▁of▁sentence｜>'}} + {%- endif %} + {%- endif %} + {%- if message['role'] == 'tool' %} + {%- set ns.is_last_user = false -%} + {%- set ns.is_tool = true -%} + {%- if ns.is_output_first %} + {{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + content + '<｜tool▁output▁end｜>'}} + {%- set ns.is_output_first = false %} + {%- else %} + {{'\n<｜tool▁output▁begin｜>' + content + '<｜tool▁output▁end｜>'}} + {%- endif %} + {%- endif %} +{%- endfor -%} +{% if ns.is_tool %} + {{'<｜tool▁outputs▁end｜>'}} +{% endif %} +{% if add_generation_prompt and not ns.is_last_user and not ns.is_tool %} + {{'<｜Assistant｜>'}} +{% endif %} diff --git a/examples/chat_template/tool_chat_template_deepseekv3.jinja b/examples/chat_template/tool_chat_template_deepseekv3.jinja new file mode 100644 index 000000000..46c1b8801 --- /dev/null +++ b/examples/chat_template/tool_chat_template_deepseekv3.jinja @@ -0,0 +1,91 @@ +{% if not add_generation_prompt is defined %} + {% set add_generation_prompt = false %} +{% endif %} + +{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true, is_last_user=false) %} +{%- for message in messages %} + {%- if message['role'] == 'system' %} + {%- if ns.is_first_sp %} + {% set ns.system_prompt = ns.system_prompt + message['content'] %} + {% set ns.is_first_sp = false %} + {%- else %} + {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %} + {%- endif %} + {%- endif %} +{%- endfor -%} + +{# --- Append tool descriptions if tools are defined --- #} +{% if tools is defined and tools is not none %} + {% set tool_ns = namespace(text='You are a helpful assistant with tool calling capabilities. ' + 'When a tool call is needed, you MUST use the following format to issue the call:\n' + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>FUNCTION_NAME\n' + '```json\n{"param1": "value1", "param2": "value2"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>\n\n' + 'Make sure the JSON is valid.' + '## Tools\n\n### Function\n\nYou have the following functions available:\n\n') %} + {% for tool in tools %} + {% set tool_ns.text = tool_ns.text + '\n```json\n' + (tool | tojson) + '\n```\n' %} + {% endfor %} + {% set ns.system_prompt = ns.system_prompt + '\n\n' + tool_ns.text %} +{% endif %} + +{{- bos_token }} +{{- ns.system_prompt }} + +{%- for message in messages %} + {%- if message['role'] == 'user' %} + {%- set ns.is_tool = false -%} + {%- set ns.is_first = false -%} + {%- set ns.is_last_user = true -%} + {{'<｜User｜>' + message['content'] + '<｜Assistant｜>'}} + {%- endif %} + {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %} + {%- set ns.is_last_user = false -%} + {%- if ns.is_tool %} + {{- '<｜tool▁outputs▁end｜>'}} + {%- endif %} + {%- set ns.is_first = false %} + {%- set ns.is_tool = false -%} + {%- set ns.is_output_first = true %} + {%- for tool in message['tool_calls'] %} + {%- if not ns.is_first %} + {%- if message['content'] is none %} + {{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments']|tojson + '\n' + '```' + '<｜tool▁call▁end｜>'}} + {%- else %} + {{- message['content'] + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments']|tojson + '\n' + '```' + '<｜tool▁call▁end｜>'}} + {%- endif %} + {%- set ns.is_first = true -%} + {%- else %} + {{- '\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments']|tojson + '\n' + '```' + '<｜tool▁call▁end｜>'}} + {%- endif %} + {%- endfor %} + {{- '<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}} + {%- endif %} + {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none)%} + {%- set ns.is_last_user = false -%} + {%- if ns.is_tool %} + {{- '<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}} + {%- set ns.is_tool = false -%} + {%- else %} + {% set content = message['content'] %} + {{- content + '<｜end▁of▁sentence｜>'}} + {%- endif %} + {%- endif %} + {%- if message['role'] == 'tool' %} + {%- set ns.is_last_user = false -%} + {%- set ns.is_tool = true -%} + {%- if ns.is_output_first %} + {{- 'Use the results below to formulate an answer to the user question unless additional information is needed.' }} + {{- '<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}} + {%- set ns.is_output_first = false %} + {%- else %} + {{- '\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}} + {%- endif %} + {%- endif %} +{%- endfor -%} + +{% if ns.is_tool %} + {{- '<｜tool▁outputs▁end｜>'}} +{% endif %} +{% if add_generation_prompt and not ns.is_last_user and not ns.is_tool %} + {{- '<｜Assistant｜>'}} +{% endif %} diff --git a/examples/chat_template/tool_chat_template_deepseekv31.jinja b/examples/chat_template/tool_chat_template_deepseekv31.jinja new file mode 100644 index 000000000..08e93a30a --- /dev/null +++ b/examples/chat_template/tool_chat_template_deepseekv31.jinja @@ -0,0 +1,91 @@ +{% if not add_generation_prompt is defined %} + {% set add_generation_prompt = false %} +{% endif %} +{% if not thinking is defined %} + {% set thinking = false %} +{% endif %} +{% set ns = namespace(is_first=false, is_tool=false, system_prompt='', is_first_sp=true, is_last_user=false) %} +{%- for message in messages %} + {%- if message['role'] == 'system' %} + {%- if ns.is_first_sp %} + {% set ns.system_prompt = ns.system_prompt + message['content'] %} + {% set ns.is_first_sp = false %} + {%- else %} + {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %} + {%- endif %} + {%- endif %} +{%- endfor %} + +{% if tools is defined and tools is not none %} + {% set tool_ns = namespace(text='## Tools\nYou have access to the following tools:\n') %} + {% for tool in tools %} + {% set tool_ns.text = tool_ns.text + '\n### ' + tool.function.name + '\nDescription: ' + tool.function.description + '\n\nParameters: ' + (tool.function.parameters | tojson) + '\n' %} + {% endfor %} + {% set tool_ns.text = tool_ns.text + "\nIMPORTANT: ALWAYS adhere to this exact format for tool use:\n<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>tool_call_name<｜tool▁sep｜>tool_call_arguments<｜tool▁call▁end｜>{{additional_tool_calls}}<｜tool▁calls▁end｜>\n\nWhere:\n\n- `tool_call_name` must be an exact match to one of the available tools\n- `tool_call_arguments` must be valid JSON that strictly follows the tool's Parameters Schema\n- For multiple tool calls, chain them directly without separators or spaces\n" %} + {% set ns.system_prompt = ns.system_prompt + '\n\n' + tool_ns.text %} +{% endif %} + +{{ bos_token }}{{ ns.system_prompt }} +{%- for message in messages %} + {%- if message['role'] == 'user' %} + {%- set ns.is_tool = false -%} + {%- set ns.is_first = false -%} + {%- set ns.is_last_user = true -%} + {{'<｜User｜>' + message['content']}} + {%- endif %} + {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %} + {%- if ns.is_last_user %} + {{'<｜Assistant｜>'}} + {%- endif %} + {%- set ns.is_last_user = false -%} + {%- set ns.is_first = false %} + {%- set ns.is_tool = false -%} + {%- for tool in message['tool_calls'] %} + {%- if not ns.is_first %} + {%- if message['content'] is none %} + {{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>'+ tool['function']['name'] + '<｜tool▁sep｜>' + tool['function']['arguments']|tojson + '<｜tool▁call▁end｜>'}} + {%- else %} + {{message['content'] + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['function']['name'] + '<｜tool▁sep｜>' + tool['function']['arguments']|tojson + '<｜tool▁call▁end｜>'}} + {%- endif %} + {%- set ns.is_first = true -%} + {%- else %} + {{'<｜tool▁call▁begin｜>'+ tool['function']['name'] + '<｜tool▁sep｜>' + tool['function']['arguments']|tojson + '<｜tool▁call▁end｜>'}} + {%- endif %} + {%- endfor %} + {{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}} + {%- endif %} + {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none) %} + {%- if ns.is_last_user %} + {{'<｜Assistant｜>'}} + {%- if message['prefix'] is defined and message['prefix'] and thinking %} + {{''}} + {%- else %} + {{''}} + {%- endif %} + {%- endif %} + {%- set ns.is_last_user = false -%} + {%- if ns.is_tool %} + {{message['content'] + '<｜end▁of▁sentence｜>'}} + {%- set ns.is_tool = false -%} + {%- else %} + {%- set content = message['content'] -%} + {%- if '' in content %} + {%- set content = content.split('', 1)[1] -%} + {%- endif %} + {{content + '<｜end▁of▁sentence｜>'}} + {%- endif %} + {%- endif %} + {%- if message['role'] == 'tool' %} + {%- set ns.is_last_user = false -%} + {%- set ns.is_tool = true -%} + {{'<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}} + {%- endif %} +{%- endfor -%} +{%- if add_generation_prompt and ns.is_last_user and not ns.is_tool %} + {{'<｜Assistant｜>'}} + {%- if not thinking %} + {{''}} + {%- else %} + {{''}} + {%- endif %} +{% endif %} diff --git a/examples/chat_template/tool_chat_template_llama4_pythonic.jinja b/examples/chat_template/tool_chat_template_llama4_pythonic.jinja new file mode 100644 index 000000000..74b315c31 --- /dev/null +++ b/examples/chat_template/tool_chat_template_llama4_pythonic.jinja @@ -0,0 +1,112 @@ +{# Copied from https://github.com/wukaixingxp/vllm/blob/8a32e2a6e452a03c0e8222e3876ad6086cbf581f/examples/tool_chat_template_llama4_pythonic.jinja to enable better model response. #} +{{- bos_token }} +{%- if custom_tools is defined and custom_tools %} + {%- set tools = custom_tools %} +{%- endif %} +{%- if tools is defined and tools %} + {%- set tool_definition = tool_definition ~ (tools | tojson(indent=4)) %} +{%- else %} + {%- set tools = none %} +{%- endif %} + + +{#- This block extracts the system message, so we can slot it into the right place. #} +{%- if messages[0]['role'] == 'system' %} + {%- set user_provided_system_message = true %} + {%- if messages[0]['content'] is string %} + {%- set system_message = messages[0]['content']|trim %} + {%- else %} + {%- set system_message = messages[0]['content'][0]['text']|trim %} + {%- endif %} + {%- set messages = messages[1:] %} +{%- else %} + {%- if tools is not none %} + {#- Since not system_message was provided by user, if tool is provided, system_message is now default tool system message #} + {#- This system message is from llama website:https://www.llama.com/docs/model-cards-and-prompt-formats/llama4/ #} + {%- set system_message = "You are a helpful assistant and an expert in function composition. You can answer general questions using your internal knowledge OR invoke functions when necessary. Follow these strict guidelines:\n\n1. FUNCTION CALLS:\n- ONLY use functions that are EXPLICITLY listed in the function list below\n- If NO functions are listed (empty function list []), respond ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\"\n- If a function is not in the list, respond ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\"\n- If ALL required parameters are present AND the query EXACTLY matches a listed function's purpose: output ONLY the function call(s)\n- Use exact format: [func_name1(param1=value1, param2=value2), func_name2(...)]\nExamples:\nCORRECT: [get_weather(location=\"Vancouver\"), calculate_route(start=\"Boston\", end=\"New York\")] <- Only if get_weather and calculate_route are in function list\nINCORRECT: get_weather(location=\"New York\")\nINCORRECT: Let me check the weather: [get_weather(location=\"New York\")]\nINCORRECT: [get_events(location=\"Singapore\")] <- If function not in list\n\n2. RESPONSE RULES:\n- For pure function requests matching a listed function: ONLY output the function call(s)\n- For knowledge questions: ONLY output text\n- For missing parameters: ONLY request the specific missing parameters\n- For unavailable services (not in function list): output ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\". Do NOT execute a function call.\n- If the query asks for information beyond what a listed function provides: output ONLY with internal knowledge about your limitations\n- NEVER combine text and function calls in the same response\n- NEVER suggest alternative functions when the requested service is unavailable\n- NEVER create or invent new functions not listed below\n\n3. STRICT BOUNDARIES:\n- ONLY use functions from the list below - no exceptions\n- NEVER use a function as an alternative to unavailable information\n- NEVER call functions not present in the function list\n- NEVER add explanatory text to function calls\n- NEVER respond with empty brackets\n- Use proper Python/JSON syntax for function calls\n- Check the function list carefully before responding\n\n4. TOOL RESPONSE HANDLING:\n- When receiving tool responses: provide concise, natural language responses\n- Don't repeat tool response verbatim\n- Don't add supplementary information\n\nHere is a list of functions in JSON format that you can invoke:\n" %} + {%- else %} + {%- set system_message = "" %} + {%- endif %} +{%- endif %} +{#- Now writing the system message: use the user provided system message if user_provided_system_message, else default tool system message if tools presented #} +{%- if system_message %} + {#- always use user provided system message to override default tool system message #} + {{- "<|header_start|>system<|header_end|>\n\n" }} + {{- system_message }} + {%- if user_provided_system_message and tools %} + {{- "\nHere is a list of functions in JSON format that you can invoke. Use exact format: [func_name1(param1=value1, param2=value2), func_name2(...)]\n" }} + {{- tool_definition -}} + {%- elif tool_definition %} + {{- tool_definition -}} + {%- endif %} + {{- "<|eot|>" }} +{%- endif %} + +{#- Now deal with all other messages #} +{%- for message in messages %} + {#- Base case: messages that are not from tool role and has empty tool_call list #} + {%- if not (message.role == 'ipython' or message.role == 'tool' or ('tool_calls' in message and message.tool_calls|length != 0 )) %} + {{- '<|header_start|>' + message['role'] + '<|header_end|>\n\n' }} + {%- if message['content'] is string %} + {{- message['content'] }} + {%- else %} + {%- for content in message['content'] %} + {%- if content['type'] == 'image' %} + {{- '<|image|>' }} + {%- elif content['type'] == 'text' %} + {{- content['text'] | trim }} + {%- endif %} + {%- endfor %} + {%- endif %} + {{- "<|eot|>" }} + {#- Tool case: messages has non-empty tool_call list, must from assistant #} + {%- elif 'tool_calls' in message %} + {#- assume tool_calls are always coming from assistant #} + {%- if message.role == 'assistant' %} + {{- '<|header_start|>assistant<|header_end|>\n\n' -}} + {%- if message['content'] is string %} + {{- message['content'] }} + {%- else %} + {%- for content in message['content'] %} + {%- if content['type'] == 'image' %} + {{- '<|image|>' }} + {%- elif content['type'] == 'text' %} + {{- content['text'] }} + {%- endif %} + {%- endfor %} + {%- endif %} + {{- "[" }} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {{- tool_call.name + '(' -}} + {%- for param in tool_call.arguments %} + {{- param + '="' -}} + {{- "%s" | format(tool_call.arguments[param]) -}} + {{- '"' -}} + {% if not loop.last %}, {% endif %} + {%- endfor %} + {{- ')' -}} + {% if not loop.last %}, {% endif %} + {%- endfor %} + {{- "]<|eot|>" }} +{%- endif %} +{#- Tool_response case: messages are from tool_response #} + {%- elif message.role == "tool" or message.role == "ipython" %} + {{- "<|header_start|>ipython<|header_end|>\n\n" }} + {%- if message.content is string %} + {{- message.content | tojson }} + {%- else %} + {%- for content in message['content'] %} + {%- if content['type'] == 'text' %} + {{- content['text'] | tojson }} + {%- endif %} + {%- endfor %} + {%- endif %} + {{- "<|eot|>" }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|header_start|>assistant<|header_end|>\n\n' }} +{%- endif %} diff --git a/examples/frontend_language/quick_start/anthropic_example_chat.py b/examples/frontend_language/quick_start/anthropic_example_chat.py new file mode 100644 index 000000000..03d699be7 --- /dev/null +++ b/examples/frontend_language/quick_start/anthropic_example_chat.py @@ -0,0 +1,73 @@ +""" +Usage: +export ANTHROPIC_API_KEY=sk-****** +python3 anthropic_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.Anthropic("claude-3-haiku-20240307")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/anthropic_example_complete.py b/examples/frontend_language/quick_start/anthropic_example_complete.py new file mode 100644 index 000000000..bce2a61ea --- /dev/null +++ b/examples/frontend_language/quick_start/anthropic_example_complete.py @@ -0,0 +1,68 @@ +""" +Usage: +export ANTHROPIC_API_KEY=sk-****** +python3 anthropic_example_complete.py +""" + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """ +\n\nHuman: What is the capital of France? +\n\nAssistant: Paris +\n\nHuman: What is the capital of Germany? +\n\nAssistant: Berlin +\n\nHuman: What is the capital of Italy? +\n\nAssistant: Rome +""" + s += "\n\nHuman: " + question + "\n" + s += "\n\nAssistant:" + sgl.gen("answer", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.Anthropic("claude-3-haiku-20240307")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/azure_openai_example_chat.py b/examples/frontend_language/quick_start/azure_openai_example_chat.py new file mode 100644 index 000000000..d53f935f4 --- /dev/null +++ b/examples/frontend_language/quick_start/azure_openai_example_chat.py @@ -0,0 +1,83 @@ +""" +Usage: +export AZURE_OPENAI_API_KEY=sk-****** +python3 openai_example_chat.py +""" + +import os + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + backend = sgl.OpenAI( + model_name="azure-gpt-4", + api_version="2023-07-01-preview", + azure_endpoint="https://oai-arena-sweden.openai.azure.com/", + api_key=os.environ["AZURE_OPENAI_API_KEY"], + is_azure=True, + ) + sgl.set_default_backend(backend) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/gemini_example_chat.py b/examples/frontend_language/quick_start/gemini_example_chat.py new file mode 100644 index 000000000..0ae623109 --- /dev/null +++ b/examples/frontend_language/quick_start/gemini_example_chat.py @@ -0,0 +1,73 @@ +""" +Usage: +export GCP_PROJECT_ID=****** +python3 gemini_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.VertexAI("gemini-pro")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/gemini_example_complete.py b/examples/frontend_language/quick_start/gemini_example_complete.py new file mode 100644 index 000000000..5188bf418 --- /dev/null +++ b/examples/frontend_language/quick_start/gemini_example_complete.py @@ -0,0 +1,68 @@ +""" +Usage: +export GCP_PROJECT_ID=****** +python3 gemini_example_complete.py +""" + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """The following are questions with answers. +Q: What is the capital of France? +A: Paris +Q: What is the capital of Germany? +A: Berlin +Q: What is the capital of Italy? +A: Rome +""" + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.VertexAI("gemini-pro")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/gemini_example_multimodal_chat.py b/examples/frontend_language/quick_start/gemini_example_multimodal_chat.py new file mode 100644 index 000000000..afe0c723f --- /dev/null +++ b/examples/frontend_language/quick_start/gemini_example_multimodal_chat.py @@ -0,0 +1,30 @@ +""" +Usage: +export GCP_PROJECT_ID=****** +python3 gemini_example_multimodal_chat.py +""" + +import sglang as sgl + + +@sgl.function +def image_qa(s, image_file1, image_file2, question): + s += sgl.user(sgl.image(image_file1) + sgl.image(image_file2) + question) + s += sgl.assistant(sgl.gen("answer", max_tokens=256)) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.VertexAI("gemini-pro-vision")) + + state = image_qa.run( + image_file1="./images/cat.jpeg", + image_file2="./images/dog.jpeg", + question="Describe difference of the two images in one sentence.", + stream=True, + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + print(state["answer"]) diff --git a/examples/frontend_language/quick_start/images/cat.jpeg b/examples/frontend_language/quick_start/images/cat.jpeg new file mode 100644 index 000000000..a6a8e48c9 Binary files /dev/null and b/examples/frontend_language/quick_start/images/cat.jpeg differ diff --git a/examples/frontend_language/quick_start/images/dog.jpeg b/examples/frontend_language/quick_start/images/dog.jpeg new file mode 100644 index 000000000..dc4b40e21 Binary files /dev/null and b/examples/frontend_language/quick_start/images/dog.jpeg differ diff --git a/examples/frontend_language/quick_start/local_example_chat.py b/examples/frontend_language/quick_start/local_example_chat.py new file mode 100644 index 000000000..e1e4b62cc --- /dev/null +++ b/examples/frontend_language/quick_start/local_example_chat.py @@ -0,0 +1,75 @@ +""" +Usage: +python3 local_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-chat-hf") + sgl.set_default_backend(runtime) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() + + runtime.shutdown() diff --git a/examples/frontend_language/quick_start/local_example_complete.py b/examples/frontend_language/quick_start/local_example_complete.py new file mode 100644 index 000000000..00a451cf6 --- /dev/null +++ b/examples/frontend_language/quick_start/local_example_complete.py @@ -0,0 +1,70 @@ +""" +Usage: +python3 local_example_complete.py +""" + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """The following are questions with answers. +Q: What is the capital of France? +A: Paris +Q: What is the capital of Germany? +A: Berlin +Q: What is the capital of Italy? +A: Rome +""" + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-chat-hf") + sgl.set_default_backend(runtime) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() + + runtime.shutdown() diff --git a/examples/frontend_language/quick_start/local_example_llava_next.py b/examples/frontend_language/quick_start/local_example_llava_next.py new file mode 100644 index 000000000..c941a549e --- /dev/null +++ b/examples/frontend_language/quick_start/local_example_llava_next.py @@ -0,0 +1,78 @@ +""" +Usage: python3 local_example_llava_next.py +""" + +import sglang as sgl +from sglang.lang.chat_template import get_chat_template + + +@sgl.function +def image_qa(s, image_path, question): + s += sgl.user(sgl.image(image_path) + question) + s += sgl.assistant(sgl.gen("answer")) + + +def single(): + state = image_qa.run( + image_path="images/cat.jpeg", question="What is this?", max_new_tokens=128 + ) + print(state["answer"], "\n") + + +def stream(): + state = image_qa.run( + image_path="images/cat.jpeg", + question="What is this?", + max_new_tokens=64, + stream=True, + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = image_qa.run_batch( + [ + {"image_path": "images/cat.jpeg", "question": "What is this?"}, + {"image_path": "images/dog.jpeg", "question": "What is this?"}, + ], + max_new_tokens=128, + ) + for s in states: + print(s["answer"], "\n") + + +if __name__ == "__main__": + import multiprocessing as mp + + mp.set_start_method("spawn", force=True) + + runtime = sgl.Runtime(model_path="lmms-lab/llama3-llava-next-8b") + runtime.endpoint.chat_template = get_chat_template("llama-3-instruct-llava") + + # Or you can use the 72B model + # runtime = sgl.Runtime(model_path="lmms-lab/llava-next-72b", tp_size=8) + # runtime.endpoint.chat_template = get_chat_template("chatml-llava") + + sgl.set_default_backend(runtime) + print(f"chat template: {runtime.endpoint.chat_template.name}") + + # Or you can use API models + # sgl.set_default_backend(sgl.OpenAI("gpt-4-vision-preview")) + # sgl.set_default_backend(sgl.VertexAI("gemini-pro-vision")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() + + runtime.shutdown() diff --git a/examples/frontend_language/quick_start/openai_example_chat.py b/examples/frontend_language/quick_start/openai_example_chat.py new file mode 100644 index 000000000..9511e21cf --- /dev/null +++ b/examples/frontend_language/quick_start/openai_example_chat.py @@ -0,0 +1,74 @@ +""" +Usage: +export OPENAI_API_KEY=sk-****** +python3 openai_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/openai_example_complete.py b/examples/frontend_language/quick_start/openai_example_complete.py new file mode 100644 index 000000000..d64bcaf1c --- /dev/null +++ b/examples/frontend_language/quick_start/openai_example_complete.py @@ -0,0 +1,68 @@ +""" +Usage: +export OPENAI_API_KEY=sk-****** +python3 openai_example_complete.py +""" + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """The following are questions with answers. +Q: What is the capital of France? +A: Paris +Q: What is the capital of Germany? +A: Berlin +Q: What is the capital of Italy? +A: Rome +""" + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo-instruct")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/openai_example_n.py b/examples/frontend_language/quick_start/openai_example_n.py new file mode 100644 index 000000000..25372b9f4 --- /dev/null +++ b/examples/frontend_language/quick_start/openai_example_n.py @@ -0,0 +1,71 @@ +""" +Usage: +export OPENAI_API_KEY=sk-****** +python3 openai_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=1024, n=2)) + s += sgl.user(question_2) + s += sgl.assistant( + sgl.gen( + "answer_2", + max_tokens=1024, + ) + ) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + print("\n-- answer_2 --\n", state["answer_2"]) + assert isinstance(state["answer_1"], list) + assert len(state["answer_1"]) == 2 + assert isinstance(state["answer_2"], str) + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + print("\n-- answer_1 --\n", s["answer_1"]) + print("\n-- answer_2 --\n", s["answer_2"]) + assert isinstance(s["answer_1"], list) + assert len(s["answer_1"]) == 2 + assert isinstance(s["answer_2"], str) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.OpenAI("o1")) + + # Run a single request + print("\n========== single ==========\n") + single() + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/openai_example_o1.py b/examples/frontend_language/quick_start/openai_example_o1.py new file mode 100644 index 000000000..2e5c14002 --- /dev/null +++ b/examples/frontend_language/quick_start/openai_example_o1.py @@ -0,0 +1,57 @@ +""" +Usage: +export OPENAI_API_KEY=sk-****** +python3 openai_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=100)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2")) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.OpenAI("o1")) + + # Run a single request + print("\n========== single ==========\n") + single() + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/openrouter_example_chat.py b/examples/frontend_language/quick_start/openrouter_example_chat.py new file mode 100644 index 000000000..a0b6f15bc --- /dev/null +++ b/examples/frontend_language/quick_start/openrouter_example_chat.py @@ -0,0 +1,81 @@ +""" +Usage: +export OPENROUTER_API_KEY=sk-****** +python3 together_example_chat.py +""" + +import os + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + backend = sgl.OpenAI( + model_name="google/gemma-7b-it:free", + base_url="https://openrouter.ai/api/v1", + api_key=os.environ.get("OPENROUTER_API_KEY"), + ) + sgl.set_default_backend(backend) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/together_example_chat.py b/examples/frontend_language/quick_start/together_example_chat.py new file mode 100644 index 000000000..2d2059062 --- /dev/null +++ b/examples/frontend_language/quick_start/together_example_chat.py @@ -0,0 +1,81 @@ +""" +Usage: +export TOGETHER_API_KEY=sk-****** +python3 together_example_chat.py +""" + +import os + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + backend = sgl.OpenAI( + model_name="mistralai/Mixtral-8x7B-Instruct-v0.1", + base_url="https://api.together.xyz/v1", + api_key=os.environ.get("TOGETHER_API_KEY"), + ) + sgl.set_default_backend(backend) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/quick_start/together_example_complete.py b/examples/frontend_language/quick_start/together_example_complete.py new file mode 100644 index 000000000..d9119ed6c --- /dev/null +++ b/examples/frontend_language/quick_start/together_example_complete.py @@ -0,0 +1,76 @@ +""" +Usage: +export TOGETHER_API_KEY=sk-****** +python3 together_example_complete.py +""" + +import os + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """The following are questions with answers. +Q: What is the capital of France? +A: Paris +Q: What is the capital of Germany? +A: Berlin +Q: What is the capital of Italy? +A: Rome +""" + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + backend = sgl.OpenAI( + model_name="mistralai/Mixtral-8x7B-Instruct-v0.1", + is_chat_model=False, + base_url="https://api.together.xyz/v1", + api_key=os.environ.get("TOGETHER_API_KEY"), + ) + sgl.set_default_backend(backend) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/examples/frontend_language/usage/chinese_regex.py b/examples/frontend_language/usage/chinese_regex.py new file mode 100644 index 000000000..78e9c7e16 --- /dev/null +++ b/examples/frontend_language/usage/chinese_regex.py @@ -0,0 +1,53 @@ +import sglang as sgl + +character_regex = ( + r"""\{\n""" + + r""" "姓名": "[^"]{1,32}",\n""" + + r""" "学院": "(格兰芬多|赫奇帕奇|拉文克劳|斯莱特林)",\n""" + + r""" "血型": "(纯血|混血|麻瓜)",\n""" + + r""" "职业": "(学生|教师|傲罗|魔法部|食死徒|凤凰社成员)",\n""" + + r""" "魔杖": \{\n""" + + r""" "材质": "[^"]{1,32}",\n""" + + r""" "杖芯": "[^"]{1,32}",\n""" + + r""" "长度": [0-9]{1,2}\.[0-9]{0,2}\n""" + + r""" \},\n""" + + r""" "存活": "(存活|死亡)",\n""" + + r""" "守护神": "[^"]{1,32}",\n""" + + r""" "博格特": "[^"]{1,32}"\n""" + + r"""\}""" +) + + +@sgl.function +def character_gen(s, name): + s += name + " 是一名哈利波特系列小说中的角色。请填写以下关于这个角色的信息。" + s += """\ +这是一个例子 +{ + "姓名": "哈利波特", + "学院": "格兰芬多", + "血型": "混血", + "职业": "学生", + "魔杖": { + "材质": "冬青木", + "杖芯": "凤凰尾羽", + "长度": 11.0 + }, + "存活": "存活", + "守护神": "麋鹿", + "博格特": "摄魂怪" +} +""" + s += f"现在请你填写{name}的信息：\n" + s += sgl.gen("json_output", max_tokens=256, regex=character_regex) + + +def main(): + backend = sgl.RuntimeEndpoint("http://localhost:30000") + sgl.set_default_backend(backend) + ret = character_gen.run(name="赫敏格兰杰", temperature=0) + print(ret.text()) + + +if __name__ == "__main__": + main() diff --git a/examples/frontend_language/usage/choices_logprob.py b/examples/frontend_language/usage/choices_logprob.py new file mode 100644 index 000000000..6cd733fe9 --- /dev/null +++ b/examples/frontend_language/usage/choices_logprob.py @@ -0,0 +1,44 @@ +""" +Usage: +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +python choices_logprob.py +""" + +import sglang as sgl + + +@sgl.function +def tool_use(s, question): + s += "To answer this question: " + question + ", " + s += "I need to use a " + sgl.gen("tool", choices=["calculator", "search engine"]) + + +def main(): + # Run one case + question = "What is 5 + 5?" + state = tool_use.run(question) + print("questions:", question) + print("choice:", state["tool"]) + meta_info = state.get_meta_info("tool") + print("logprobs of choice 1", meta_info["input_token_logprobs"][0]) + print("logprobs of choice 2", meta_info["input_token_logprobs"][1]) + print("-" * 50) + + # Run a batch + questions = [ + "What is 5 + 6?", + "Who is Michael Jordan?", + ] + states = tool_use.run_batch([{"question": q} for q in questions]) + for question, state in zip(questions, states): + print("questions:", question) + print("choice:", state["tool"]) + meta_info = state.get_meta_info("tool") + print("logprobs of choice 1", meta_info["input_token_logprobs"][0]) + print("logprobs of choice 2", meta_info["input_token_logprobs"][1]) + print("-" * 50) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + main() diff --git a/examples/frontend_language/usage/cot_decoding.py b/examples/frontend_language/usage/cot_decoding.py new file mode 100644 index 000000000..7a7a04bce --- /dev/null +++ b/examples/frontend_language/usage/cot_decoding.py @@ -0,0 +1,115 @@ +from math import exp +from pprint import pformat + +import sglang as sgl + +YELLOW = "\033[1;33m" +GREEN = "\033[1;32m" +BLUE = "\033[1;34m" +CLEAR = "\033[1;0m" + + +@sgl.function +def cot_decoding(s, question, get_top_k, is_chat_model, verbose): + """CoT Decoding: http://arxiv.org/abs/2402.10200""" + + if is_chat_model: + s += sgl.user("Question: " + question + "\nAnswer:") + s += sgl.assistant_begin() + else: + s += "Question: " + question + "\nAnswer:" + + step_0 = s.fork(1)[0] + forks = s.fork(get_top_k) + answer_forks = s.fork(get_top_k) + + # decoding step 0 + step_0 += sgl.gen( + "get_top_k", + max_tokens=0, + return_logprob=True, + top_logprobs_num=get_top_k, + return_text_in_logprobs=True, + ) + logprobs = step_0.get_meta_info("get_top_k")["output_top_logprobs"][0] + + print("Decoding step 0:", ", ".join(pformat(token[2]) for token in logprobs)) + for idx, (f, token) in enumerate(zip(forks, logprobs)): + logprob, token_id, text = token + f += text + + if text == "<|end_of_text|>": + print( + f"{YELLOW}Path #{idx} {pformat(text)}[{exp(logprob):.3f}] (score=nan, answer=nan){CLEAR}" + ) + continue + + # continue greedy decoding + f += sgl.gen( + "answer", + temperature=0, + max_tokens=1024, + return_logprob=True, + top_logprobs_num=2, + return_text_in_logprobs=True, + ) + + # calculate probability disparity between the top and secondary tokens + x1s = [exp(xt[0][0]) for xt in f.get_meta_info("answer")["output_top_logprobs"]] + x2s = [exp(xt[1][0]) for xt in f.get_meta_info("answer")["output_top_logprobs"]] + tokens = [xt[0][2] for xt in f.get_meta_info("answer")["output_top_logprobs"]] + delta = (sum(x1s) - sum(x2s)) / len(x1s) + + # extract the answer span (without the '<|end_of_text|>' token) + answer_forks[idx] += text + f["answer"] + "\nSo the answer is" + answer_forks[idx] += sgl.gen( + "answer_span", + temperature=0, + max_tokens=64, + return_logprob=True, + top_logprobs_num=2, + return_text_in_logprobs=True, + ) + answer = answer_forks[idx]["answer_span"].replace("\n", " ").strip(":") + print( + f"{YELLOW}Path #{idx} {pformat(text)}[{exp(logprob):.3f}] (score={delta}, answer={answer}){CLEAR}" + ) + generated_text = str(answer_forks[idx])[len("ProgramState(") : -1] + print(f"{BLUE}{pformat(generated_text)}{CLEAR}") + + if verbose: + answer_tokens = [ + xt[0][2] + for xt in answer_forks[idx].get_meta_info("answer_span")[ + "output_top_logprobs" + ] + ] + answer_x1s = [ + exp(xt[0][0]) + for xt in answer_forks[idx].get_meta_info("answer_span")[ + "output_top_logprobs" + ] + ] + answer_x2s = [ + exp(xt[1][0]) + for xt in answer_forks[idx].get_meta_info("answer_span")[ + "output_top_logprobs" + ] + ] + + for token, x1, x2 in zip(tokens, x1s, x2s): + print(f" {GREEN}{pformat(token)}{CLEAR}({x1:.3f}-{x2:.3f})", end="") + print("\n===========") + for token, x1, x2 in zip(answer_tokens, answer_x1s, answer_x2s): + print(f" {GREEN}{pformat(token)}{CLEAR}({x1:.3f}-{x2:.3f})", end="") + print() + + +sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + +state = cot_decoding.run( + question=r"Claire makes a 3 egg omelet every morning for breakfast. How many dozens of eggs will she eat in 4 weeks?", + get_top_k=10, + is_chat_model=True, + verbose=False, +) diff --git a/examples/frontend_language/usage/json_decode.py b/examples/frontend_language/usage/json_decode.py new file mode 100644 index 000000000..5dc3522d5 --- /dev/null +++ b/examples/frontend_language/usage/json_decode.py @@ -0,0 +1,83 @@ +""" +Usage: +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +python json_decode.py +""" + +from enum import Enum + +from pydantic import BaseModel + +import sglang as sgl +from sglang.srt.constrained.outlines_backend import build_regex_from_object + +character_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "house": "(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)",\n""" + + r""" "blood status": "(Pure-blood|Half-blood|Muggle-born)",\n""" + + r""" "occupation": "(student|teacher|auror|ministry of magic|death eater|order of the phoenix)",\n""" + + r""" "wand": \{\n""" + + r""" "wood": "[\w\d\s]{1,16}",\n""" + + r""" "core": "[\w\d\s]{1,16}",\n""" + + r""" "length": [0-9]{1,2}\.[0-9]{0,2}\n""" + + r""" \},\n""" + + r""" "alive": "(Alive|Deceased)",\n""" + + r""" "patronus": "[\w\d\s]{1,16}",\n""" + + r""" "bogart": "[\w\d\s]{1,16}"\n""" + + r"""\}""" +) + + +@sgl.function +def character_gen(s, name): + s += ( + name + + " is a character in Harry Potter. Please fill in the following information about this character.\n" + ) + s += "The constrained regex is:\n" + s += character_regex + "\n" + s += "The JSON output is:\n" + s += sgl.gen("json_output", max_tokens=256, regex=character_regex) + + +def driver_character_gen(): + state = character_gen.run(name="Hermione Granger") + print(state.text()) + + +class Weapon(str, Enum): + sword = "sword" + axe = "axe" + mace = "mace" + spear = "spear" + bow = "bow" + crossbow = "crossbow" + + +class Wizard(BaseModel): + name: str + age: int + weapon: Weapon + + +@sgl.function +def pydantic_wizard_gen(s): + s += "Give me a description about a wizard in the JSON format.\n" + s += sgl.gen( + "character", + max_tokens=128, + temperature=0, + regex=build_regex_from_object(Wizard), # Requires pydantic >= 2.0 + ) + + +def driver_pydantic_wizard_gen(): + state = pydantic_wizard_gen.run() + print(state.text()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + driver_character_gen() + # driver_pydantic_wizard_gen() diff --git a/examples/frontend_language/usage/json_logprobs.py b/examples/frontend_language/usage/json_logprobs.py new file mode 100644 index 000000000..15206a619 --- /dev/null +++ b/examples/frontend_language/usage/json_logprobs.py @@ -0,0 +1,103 @@ +# NOTE: Currently this can only be run through HTTP requests. +from concurrent.futures import ThreadPoolExecutor + +from json_decode import character_regex + +from sglang.utils import http_request + +character_names = ["Hermione Granger", "Ron Weasley", "Harry Potter"] + +base_url = "http://localhost:30000" + +prompt = "is a character in Harry Potter. Please fill in the following information about this character.\n" + + +def openai_api_request(name): + data = { + "model": "", + "prompt": name + prompt, + "temperature": 0, + "max_tokens": 128, + "regex": character_regex, + "logprobs": 3, + } + res = http_request(base_url + "/v1/completions", json=data).json() + + # with open(f"json_logprobs_{name.replace(' ', '_')}_tmp.json", "w") as fout: + # fout.write(json.dumps(res, indent=4)) + + logprobs = res["choices"][0]["logprobs"] + usage = res["usage"] + assert len(logprobs["token_logprobs"]) == len(logprobs["tokens"]) + assert len(logprobs["token_logprobs"]) == len(logprobs["top_logprobs"]) + assert len(logprobs["token_logprobs"]) == usage["completion_tokens"] - 1 + + return res + + +def srt_api_request(name): + data = { + "text": name + prompt, + "sampling_params": { + "temperature": 0, + "max_new_tokens": 128, + "regex": character_regex, + }, + "return_logprob": True, + "logprob_start_len": 0, + "top_logprobs_num": 3, + "return_text_in_logprobs": True, + } + + res = http_request(base_url + "/generate", json=data).json() + + # with open(f"json_logprobs_{name.replace(' ', '_')}_tmp.json", "w") as fout: + # fout.write(json.dumps(res, indent=4)) + + meta_info = res["meta_info"] + assert len(meta_info["input_token_logprobs"]) == len( + meta_info["input_top_logprobs"] + ) + assert len(meta_info["output_token_logprobs"]) == len( + meta_info["output_top_logprobs"] + ) + assert len(meta_info["input_token_logprobs"]) == meta_info["prompt_tokens"] + assert len(meta_info["output_token_logprobs"]) == meta_info["completion_tokens"] - 1 + + return res + + +def pretty_print(res): + meta_info = res["meta_info"] + + print("\n\n", "=" * 30, "Prefill", "=" * 30) + for i in range(len(meta_info["input_token_logprobs"])): + print(f"{str(meta_info['input_token_logprobs'][i][2].encode()): <20}", end="") + top_ks = ( + [str(t[2].encode()) for t in meta_info["input_top_logprobs"][i]] + if meta_info["input_top_logprobs"][i] + else [] + ) + for top_k in top_ks: + print(f"{top_k: <15}", end="") + print() + + print("\n\n", "=" * 30, "Decode", "=" * 30) + for i in range(len(meta_info["output_token_logprobs"])): + print(f"{str(meta_info['output_token_logprobs'][i][2].encode()): <20}", end="") + top_ks = [str(t[2].encode()) for t in meta_info["output_top_logprobs"][i]] + for top_k in top_ks: + print(f"{top_k: <15}", end="") + print() + + print(res["text"]) + + +if __name__ == "__main__": + with ThreadPoolExecutor() as executor: + ress = executor.map(srt_api_request, character_names) + + for res in ress: + pretty_print(res) + + openai_api_request("Hermione Granger") diff --git a/examples/frontend_language/usage/llava_video/srt_example_llava_v.py b/examples/frontend_language/usage/llava_video/srt_example_llava_v.py new file mode 100644 index 000000000..ec5b334b0 --- /dev/null +++ b/examples/frontend_language/usage/llava_video/srt_example_llava_v.py @@ -0,0 +1,260 @@ +""" +Usage: +pip install opencv-python-headless + +python3 srt_example_llava_v.py +""" + +import argparse +import csv +import json +import os +import time + +import requests + +import sglang as sgl + + +@sgl.function +def video_qa(s, num_frames, video_path, question): + s += sgl.user(sgl.video(video_path, num_frames) + question) + s += sgl.assistant(sgl.gen("answer")) + + +def single(path, num_frames=16): + state = video_qa.run( + num_frames=num_frames, + video_path=path, + question="Please provide a detailed description of the video, focusing on the main subjects, their actions, the background scenes", + temperature=0.0, + max_new_tokens=1024, + ) + print(state["answer"], "\n") + + +def split_into_chunks(lst, num_chunks): + """Split a list into a specified number of chunks.""" + # Calculate the chunk size using integer division. Note that this may drop some items if not evenly divisible. + chunk_size = len(lst) // num_chunks + + if chunk_size == 0: + chunk_size = len(lst) + # Use list comprehension to generate chunks. The last chunk will take any remainder if the list size isn't evenly divisible. + chunks = [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)] + # Ensure we have exactly num_chunks chunks, even if some are empty + chunks.extend([[] for _ in range(num_chunks - len(chunks))]) + return chunks + + +def save_batch_results(batch_video_files, states, cur_chunk, batch_idx, save_dir): + csv_filename = f"{save_dir}/chunk_{cur_chunk}_batch_{batch_idx}.csv" + with open(csv_filename, "w", newline="") as csvfile: + writer = csv.writer(csvfile) + writer.writerow(["video_name", "answer"]) + for video_path, state in zip(batch_video_files, states): + video_name = os.path.basename(video_path) + writer.writerow([video_name, state["answer"]]) + + +def compile_and_cleanup_final_results(cur_chunk, num_batches, save_dir): + final_csv_filename = f"{save_dir}/final_results_chunk_{cur_chunk}.csv" + with open(final_csv_filename, "w", newline="") as final_csvfile: + writer = csv.writer(final_csvfile) + writer.writerow(["video_name", "answer"]) + for batch_idx in range(num_batches): + batch_csv_filename = f"{save_dir}/chunk_{cur_chunk}_batch_{batch_idx}.csv" + with open(batch_csv_filename, "r") as batch_csvfile: + reader = csv.reader(batch_csvfile) + next(reader) # Skip header row + for row in reader: + writer.writerow(row) + os.remove(batch_csv_filename) + + +def find_video_files(video_dir): + # Check if the video_dir is actually a file + if os.path.isfile(video_dir): + # If it's a file, return it as a single-element list + return [video_dir] + + # Original logic to find video files in a directory + video_files = [] + for root, dirs, files in os.walk(video_dir): + for file in files: + if file.endswith((".mp4", ".avi", ".mov")): + video_files.append(os.path.join(root, file)) + return video_files + + +def batch(video_dir, save_dir, cur_chunk, num_chunks, num_frames=16, batch_size=64): + video_files = find_video_files(video_dir) + chunked_video_files = split_into_chunks(video_files, num_chunks)[cur_chunk] + num_batches = 0 + + for i in range(0, len(chunked_video_files), batch_size): + batch_video_files = chunked_video_files[i : i + batch_size] + print(f"Processing batch of {len(batch_video_files)} video(s)...") + + if not batch_video_files: + print("No video files found in the specified directory.") + return + + batch_input = [ + { + "num_frames": num_frames, + "video_path": video_path, + "question": "Please provide a detailed description of the video, focusing on the main subjects, their actions, the background scenes.", + } + for video_path in batch_video_files + ] + + start_time = time.perf_counter() + states = video_qa.run_batch(batch_input, max_new_tokens=512, temperature=0.2) + total_time = time.perf_counter() - start_time + average_time = total_time / len(batch_video_files) + print( + f"Number of videos in batch: {len(batch_video_files)}. Average processing time per video: {average_time:.2f} seconds. Total time for this batch: {total_time:.2f} seconds" + ) + + save_batch_results(batch_video_files, states, cur_chunk, num_batches, save_dir) + num_batches += 1 + + compile_and_cleanup_final_results(cur_chunk, num_batches, save_dir) + + +if __name__ == "__main__": + + url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4" + + cache_dir = os.path.expanduser("~/.cache") + file_path = os.path.join(cache_dir, "jobs.mp4") + + os.makedirs(cache_dir, exist_ok=True) + + response = requests.get(url) + response.raise_for_status() # Raise an exception for bad responses + + with open(file_path, "wb") as f: + f.write(response.content) + + print(f"File downloaded and saved to: {file_path}") + # Create the parser + parser = argparse.ArgumentParser( + description="Run video processing with specified port." + ) + + # Add an argument for the port + parser.add_argument( + "--port", + type=int, + default=30000, + help="The master port for distributed serving.", + ) + parser.add_argument( + "--chunk-idx", type=int, default=0, help="The index of the chunk to process." + ) + parser.add_argument( + "--num-chunks", type=int, default=8, help="The number of chunks to process." + ) + parser.add_argument( + "--save-dir", + type=str, + default="./work_dirs/llava_video", + help="The directory to save the processed video files.", + ) + parser.add_argument( + "--video-dir", + type=str, + default=os.path.expanduser("~/.cache/jobs.mp4"), + help="The directory or path for the processed video files.", + ) + parser.add_argument( + "--model-path", + type=str, + default="lmms-lab/LLaVA-NeXT-Video-7B", + help="The model path for the video processing.", + ) + parser.add_argument( + "--num-frames", + type=int, + default=16, + help="The number of frames to process in each video.", + ) + parser.add_argument("--mm_spatial_pool_stride", type=int, default=2) + + # Parse the arguments + args = parser.parse_args() + cur_port = args.port + cur_chunk = args.chunk_idx + num_chunks = args.num_chunks + num_frames = args.num_frames + + if "34b" in args.model_path.lower(): + tokenizer_path = "liuhaotian/llava-v1.6-34b-tokenizer" + elif "7b" in args.model_path.lower(): + tokenizer_path = "llava-hf/llava-1.5-7b-hf" + else: + print("Invalid model path. Please specify a valid model path.") + exit() + + model_override_args = {} + model_override_args["mm_spatial_pool_stride"] = args.mm_spatial_pool_stride + model_override_args["architectures"] = ["LlavaVidForCausalLM"] + model_override_args["num_frames"] = args.num_frames + model_override_args["model_type"] = "llava" + + if "34b" in args.model_path.lower(): + model_override_args["image_token_index"] = 64002 + + if args.num_frames == 32: + model_override_args["rope_scaling"] = {"factor": 2.0, "rope_type": "linear"} + model_override_args["max_sequence_length"] = 4096 * 2 + model_override_args["tokenizer_model_max_length"] = 4096 * 2 + elif args.num_frames < 32: + pass + else: + print( + "The maximum number of frames to process is 32. Please specify a valid number of frames." + ) + exit() + + runtime = sgl.Runtime( + model_path=args.model_path, # "liuhaotian/llava-v1.6-vicuna-7b", + tokenizer_path=tokenizer_path, + port=cur_port, + json_model_override_args=json.dumps(model_override_args), + tp_size=1, + ) + sgl.set_default_backend(runtime) + print(f"chat template: {runtime.endpoint.chat_template.name}") + + # Run a single request + print("\n========== single ==========\n") + root = args.video_dir + if os.path.isfile(root): + video_files = [root] + else: + video_files = [ + os.path.join(root, f) + for f in os.listdir(root) + if f.endswith((".mp4", ".avi", ".mov")) + ] # Add more extensions if needed + start_time = time.perf_counter() # Start time for processing a single video + for cur_video in video_files[:1]: + print(cur_video) + single(cur_video, num_frames) + end_time = time.perf_counter() # End time for processing a single video + total_time = end_time - start_time + average_time = total_time / len( + video_files + ) # Calculate the average processing time + print(f"Average processing time per video: {average_time:.2f} seconds") + runtime.shutdown() + + # # Run a batch of requests + # print("\n========== batch ==========\n") + # if not os.path.exists(args.save_dir): + # os.makedirs(args.save_dir) + # batch(args.video_dir, args.save_dir, cur_chunk, num_chunks, num_frames, num_chunks) + # runtime.shutdown() diff --git a/examples/frontend_language/usage/llava_video/srt_example_llava_v.sh b/examples/frontend_language/usage/llava_video/srt_example_llava_v.sh new file mode 100755 index 000000000..ffb1af96d --- /dev/null +++ b/examples/frontend_language/usage/llava_video/srt_example_llava_v.sh @@ -0,0 +1,131 @@ +#!/bin/bash + +##### USAGE ##### +# - First node: +# ```sh +# bash examples/usage/llava_video/srt_example_llava_v.sh K 0 YOUR_VIDEO_PATH YOUR_MODEL_PATH FRAMES_PER_VIDEO +# ``` +# - Second node: +# ```sh +# bash examples/usage/llava_video/srt_example_llava_v.sh K 1 YOUR_VIDEO_PATH YOUR_MODEL_PATH FRAMES_PER_VIDEO +# ``` +# - The K node: +# ```sh +# bash examples/usage/llava_video/srt_example_llava_v.sh K K-1 YOUR_VIDEO_PATH YOUR_MODEL_PATH FRAMES_PER_VIDEO +# ``` + + +# Replace `K`, `YOUR_VIDEO_PATH`, `YOUR_MODEL_PATH`, and `FRAMES_PER_VIDEO` with your specific details. +# CURRENT_ROOT="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" +CURRENT_ROOT=$(dirname "$0") + +echo ${CURRENT_ROOT} + +cd ${CURRENT_ROOT} + +export PYTHONWARNINGS=ignore + +START_TIME=$(date +%s) # Capture start time + +NUM_NODES=$1 + +CUR_NODES_IDX=$2 + +VIDEO_DIR=$3 + +MODEL_PATH=$4 + +NUM_FRAMES=$5 + + +# FRAME_FORMAT=$6 + +# FRAME_FORMAT=$(echo $FRAME_FORMAT | tr '[:lower:]' '[:upper:]') + +# # Check if FRAME_FORMAT is either JPEG or PNG +# if [[ "$FRAME_FORMAT" != "JPEG" && "$FRAME_FORMAT" != "PNG" ]]; then +# echo "Error: FRAME_FORMAT must be either JPEG or PNG." +# exit 1 +# fi + +# export TARGET_FRAMES=$TARGET_FRAMES + +echo "Each video you will sample $NUM_FRAMES frames" + +# export FRAME_FORMAT=$FRAME_FORMAT + +# echo "The frame format is $FRAME_FORMAT" + +# Assuming GPULIST is a bash array containing your GPUs +GPULIST=(0 1 2 3 4 5 6 7) +LOCAL_CHUNKS=${#GPULIST[@]} + +echo "Number of GPUs in GPULIST: $LOCAL_CHUNKS" + +ALL_CHUNKS=$((NUM_NODES * LOCAL_CHUNKS)) + +# Calculate GPUs per chunk +GPUS_PER_CHUNK=1 + +echo $GPUS_PER_CHUNK + +for IDX in $(seq 1 $LOCAL_CHUNKS); do + ( + START=$(((IDX-1) * GPUS_PER_CHUNK)) + LENGTH=$GPUS_PER_CHUNK # Length for slicing, not the end index + + CHUNK_GPUS=(${GPULIST[@]:$START:$LENGTH}) + + # Convert the chunk GPUs array to a comma-separated string + CHUNK_GPUS_STR=$(IFS=,; echo "${CHUNK_GPUS[*]}") + + LOCAL_IDX=$((CUR_NODES_IDX * LOCAL_CHUNKS + IDX)) + + echo "Chunk $(($LOCAL_IDX - 1)) will run on GPUs $CHUNK_GPUS_STR" + + # Calculate the port for this chunk. Ensure it's incremented by 5 for each chunk. + PORT=$((10000 + RANDOM % 55536)) + + MAX_RETRIES=10 + RETRY_COUNT=0 + COMMAND_STATUS=1 # Initialize as failed + + while [ $RETRY_COUNT -lt $MAX_RETRIES ] && [ $COMMAND_STATUS -ne 0 ]; do + echo "Running chunk $(($LOCAL_IDX - 1)) on GPUs $CHUNK_GPUS_STR with port $PORT. Attempt $(($RETRY_COUNT + 1))" + +#!/bin/bash + CUDA_VISIBLE_DEVICES=$CHUNK_GPUS_STR python3 srt_example_llava_v.py \ + --port $PORT \ + --num-chunks $ALL_CHUNKS \ + --chunk-idx $(($LOCAL_IDX - 1)) \ + --save-dir work_dirs/llava_next_video_inference_results \ + --video-dir $VIDEO_DIR \ + --model-path $MODEL_PATH \ + --num-frames $NUM_FRAMES #& + + wait $! # Wait for the process to finish and capture its exit status + COMMAND_STATUS=$? + + if [ $COMMAND_STATUS -ne 0 ]; then + echo "Execution failed for chunk $(($LOCAL_IDX - 1)), attempt $(($RETRY_COUNT + 1)). Retrying..." + RETRY_COUNT=$(($RETRY_COUNT + 1)) + sleep 180 # Wait a bit before retrying + else + echo "Execution succeeded for chunk $(($LOCAL_IDX - 1))." + fi + done + + if [ $COMMAND_STATUS -ne 0 ]; then + echo "Execution failed for chunk $(($LOCAL_IDX - 1)) after $MAX_RETRIES attempts." + fi + ) #& + sleep 2 # Slight delay to stagger the start times +done + +wait + +cat work_dirs/llava_next_video_inference_results/final_results_chunk_*.csv > work_dirs/llava_next_video_inference_results/final_results_node_${CUR_NODES_IDX}.csv + +END_TIME=$(date +%s) # Capture end time +ELAPSED_TIME=$(($END_TIME - $START_TIME)) +echo "Total execution time: $ELAPSED_TIME seconds." diff --git a/examples/frontend_language/usage/openai_chat_speculative.py b/examples/frontend_language/usage/openai_chat_speculative.py new file mode 100644 index 000000000..f3fd74ed8 --- /dev/null +++ b/examples/frontend_language/usage/openai_chat_speculative.py @@ -0,0 +1,155 @@ +""" +Usage: +***Note: for speculative execution to work, user must put all "gen" in "assistant". +Show in "assistant" the desired answer format. Each "gen" term should have a stop token. +The stream mode is not supported in speculative execution. + +E.g. +correct: + sgl.assistant("\nName:" + sgl.gen("name", stop="\n") + "\nBirthday:" + sgl.gen("birthday", stop="\n") + "\nJob:" + sgl.gen("job", stop="\n")) +incorrect: + s += sgl.assistant("\nName:" + sgl.gen("name", stop="\n")) + s += sgl.assistant("\nBirthday:" + sgl.gen("birthday", stop="\n")) + s += sgl.assistant("\nJob:" + sgl.gen("job", stop="\n")) + +export OPENAI_API_KEY=sk-****** +python3 openai_chat_speculative.py +""" + +import sglang as sgl +from sglang import OpenAI, function, set_default_backend + + +@function(num_api_spec_tokens=256) +def gen_character_spec(s): + s += sgl.system("You are a helpful assistant.") + s += sgl.user("Construct a character within the following format:") + s += sgl.assistant( + "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n" + ) + s += sgl.user("Please generate new Name, Birthday and Job.\n") + s += sgl.assistant( + "Name:" + + sgl.gen("name", stop="\n") + + "\nBirthday:" + + sgl.gen("birthday", stop="\n") + + "\nJob:" + + sgl.gen("job", stop="\n") + ) + + +@function(num_api_spec_tokens=256) +def gen_character_spec_no_few_shot(s): + s += sgl.user("Construct a character. For each field stop with a newline\n") + s += sgl.assistant( + "Name:" + + sgl.gen("name", stop="\n") + + "\nAge:" + + sgl.gen("age", stop="\n") + + "\nJob:" + + sgl.gen("job", stop="\n") + ) + + +@function +def gen_character_normal(s): + s += sgl.system("You are a helpful assistant.") + s += sgl.user("What's the answer of 23 + 8?") + s += sgl.assistant(sgl.gen("answer", max_tokens=64)) + + +@function(num_api_spec_tokens=1024) +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user("Answer questions in the following format:") + s += sgl.user( + "Question 1: What is the capital of France?\nQuestion 2: What is the population of this city?\n" + ) + s += sgl.assistant( + "Answer 1: The capital of France is Paris.\nAnswer 2: The population of Paris in 2024 is estimated to be around 2.1 million for the city proper.\n" + ) + s += sgl.user("Question 1: " + question_1 + "\nQuestion 2: " + question_2) + s += sgl.assistant( + "Answer 1: " + + sgl.gen("answer_1", stop="\n") + + "\nAnswer 2: " + + sgl.gen("answer_2", stop="\n") + ) + + +def test_spec_single_turn(): + backend.token_usage.reset() + + state = gen_character_spec.run() + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- name:", state["name"]) + print("-- birthday:", state["birthday"]) + print("-- job:", state["job"]) + print(backend.token_usage) + + +def test_inaccurate_spec_single_turn(): + state = gen_character_spec_no_few_shot.run() + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- name:", state["name"]) + print("\n-- age:", state["age"]) + print("\n-- job:", state["job"]) + + +def test_normal_single_turn(): + state = gen_character_normal.run() + for m in state.messages(): + print(m["role"], ":", m["content"]) + + +def test_spec_multi_turn(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions in the capital of the United States.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + print("\n-- answer_2 --\n", state["answer_2"]) + + +def test_spec_multi_turn_stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + + +if __name__ == "__main__": + backend = OpenAI("gpt-4-turbo") + set_default_backend(backend) + + print("\n========== test spec single turn ==========\n") + # expect reasonable answer for each field + test_spec_single_turn() + + print("\n========== test inaccurate spec single turn ==========\n") + # expect incomplete or unreasonable answers + test_inaccurate_spec_single_turn() + + print("\n========== test normal single turn ==========\n") + # expect reasonable answer + test_normal_single_turn() + + print("\n========== test spec multi turn ==========\n") + # expect answer with same format as in the few shot + test_spec_multi_turn() + + print("\n========== test spec multi turn stream ==========\n") + # expect error in stream_executor: stream is not supported... + test_spec_multi_turn_stream() diff --git a/examples/frontend_language/usage/openai_speculative.py b/examples/frontend_language/usage/openai_speculative.py new file mode 100644 index 000000000..4389cb059 --- /dev/null +++ b/examples/frontend_language/usage/openai_speculative.py @@ -0,0 +1,54 @@ +""" +Usage: +python3 openai_speculative.py +""" + +from sglang import OpenAI, function, gen, set_default_backend + + +@function(num_api_spec_tokens=64) +def gen_character_spec(s): + s += "Construct a character within the following format:\n" + s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n" + s += "\nPlease generate new Name, Birthday and Job.\n" + s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n") + s += "\nJob:" + gen("job", stop="\n") + "\n" + + +@function +def gen_character_no_spec(s): + s += "Construct a character within the following format:\n" + s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n" + s += "\nPlease generate new Name, Birthday and Job.\n" + s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n") + s += "\nJob:" + gen("job", stop="\n") + "\n" + + +@function(num_api_spec_tokens=64) +def gen_character_spec_no_few_shot(s): + # s += "Construct a character with name, birthday, and job:\n" + s += "Construct a character:\n" + s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n") + s += "\nJob:" + gen("job", stop="\n") + "\n" + + +if __name__ == "__main__": + backend = OpenAI("gpt-3.5-turbo-instruct") + set_default_backend(backend) + + for function in [ + gen_character_spec, + gen_character_no_spec, + gen_character_spec_no_few_shot, + ]: + backend.token_usage.reset() + + print(f"function: {function.func.__name__}") + + state = function.run() + + print("...name:", state["name"]) + print("...birthday:", state["birthday"]) + print("...job:", state["job"]) + print(backend.token_usage) + print() diff --git a/examples/frontend_language/usage/parallel_sample.py b/examples/frontend_language/usage/parallel_sample.py new file mode 100644 index 000000000..0f3cf1700 --- /dev/null +++ b/examples/frontend_language/usage/parallel_sample.py @@ -0,0 +1,40 @@ +""" +Usage: +python3 parallel_sample.py +""" + +import sglang as sgl + + +@sgl.function +def parallel_sample(s, question, n): + s += ( + "Question: Compute 1 + 2 + 3\n" + "Reasoning: I need to use a calculator.\n" + "Tool: calculator\n" + "Answer: 6\n" + "Question: Compute 3 + 2 + 2\n" + "Reasoning: I will try a calculator.\n" + "Tool: calculator\n" + "Answer: 7\n" + ) + s += "Question: " + question + "\n" + forks = s.fork(n) + forks += "Reasoning:" + sgl.gen("reasoning", stop="\n") + "\n" + forks += "Tool:" + sgl.gen("tool", choices=["calculator", "browser"]) + "\n" + forks += "Answer:" + sgl.gen("answer", stop="\n") + "\n" + forks.join() + + +sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo-instruct")) +# sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + +state = parallel_sample.run(question="Compute 5 + 2 + 4.", n=5, temperature=1.0) + +for i in range(5): + obj = { + "reasoning": state["reasoning"][i], + "tool": state["tool"][i], + "answer": state["answer"][i], + } + print(f"[{i}], {obj}") diff --git a/examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb b/examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb new file mode 100644 index 000000000..f309142ad --- /dev/null +++ b/examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb @@ -0,0 +1,408 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# RAG Powered by SGLang & Chroma Evaluated using Parea\n", + "\n", + "In this notebook, we will build a simple RAG pipeline using SGLang to execute our LLM calls, Chroma as vector database for retrieval and [Parea](https://www.parea.ai) for tracing and evaluation. We will then evaluate the performance of our RAG pipeline. The dataset we will use was created by [Virat](https://twitter.com/virattt) and contains 100 questions, contexts and answers from the Airbnb 2023 10k filing.\n", + "\n", + "The RAG pipeline consists of two steps:\n", + "1. Retrieval: Given a question, we retrieve the relevant context from all provided contexts.\n", + "2. Generation: Given the question and the retrieved context, we generate an answer.\n", + "\n", + "ℹ️ This notebook requires an OpenAI API key.\n", + "\n", + "ℹ️ This notebook requires a Parea API key, which can be created [here](https://docs.parea.ai/api-reference/authentication#parea-api-key)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setting up the environment\n", + "\n", + "We will first install the necessary packages: `sglang`, `parea-ai` and `chromadb`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# note, if you use a Mac M1 chip, you might need to install grpcio 1.59.0 first such that installing chromadb works\n", + "# !pip install grpcio==1.59.0\n", + "\n", + "!pip install sglang[openai] parea-ai chromadb" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create a Parea API key as outlined [here](https://docs.parea.ai/api-reference/authentication#parea-api-key) and save it in a `.env` file as `PAREA_API_KEY=your-api-key`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Indexing the data\n", + "\n", + "Now it's time to download the data & index it! For that, we create a collection called `contexts` in Chroma and add the contexts as documents." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\n", + "from typing import List\n", + "\n", + "import chromadb\n", + "\n", + "path_qca = \"airbnb-2023-10k-qca.json\"\n", + "\n", + "if not os.path.exists(path_qca):\n", + " !wget https://virattt.github.io/datasets/abnb-2023-10k.json -O airbnb-2023-10k-qca.json\n", + "\n", + "with open(path_qca, \"r\") as f:\n", + " question_context_answers = json.load(f)\n", + "\n", + "chroma_client = chromadb.PersistentClient()\n", + "collection = chroma_client.get_or_create_collection(name=\"contexts\")\n", + "if collection.count() == 0:\n", + " collection.add(\n", + " documents=[qca[\"context\"] for qca in question_context_answers],\n", + " ids=[str(i) for i in range(len(question_context_answers))],\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Defining the RAG pipeline\n", + "\n", + "We will start with importing the necessary packages, setting up tracing of OpenAI calls via Parea and setting OpenAI as the default backend for SGLang." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import time\n", + "\n", + "from dotenv import load_dotenv\n", + "\n", + "from sglang import function, user, assistant, gen, set_default_backend, OpenAI\n", + "from sglang.lang.interpreter import ProgramState\n", + "from parea import Parea, trace\n", + "\n", + "\n", + "load_dotenv()\n", + "\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", + "\n", + "p = Parea(api_key=os.getenv(\"PAREA_API_KEY\"), project_name=\"rag_sglang\")\n", + "p.integrate_with_sglang()\n", + "\n", + "set_default_backend(OpenAI(\"gpt-3.5-turbo\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can define our retrieval step shown below. Notice, the `trace` decorator which will automatically trace inputs, output, latency, etc. of that call." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@trace\n", + "def retrieval(question: str) -> List[str]:\n", + " return collection.query(query_texts=[question], n_results=1)[\"documents\"][0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next we will define the generation step which uses SGLang to execute the LLM call." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def generation_sglang(s, question: str, *context: str):\n", + " context = \"\\n\".join(context)\n", + " s += user(\n", + " f\"Given this question:\\n{question}\\n\\nAnd this context:\\n{context}\\n\\nAnswer the question.\"\n", + " )\n", + " s += assistant(gen(\"answer\"))\n", + "\n", + "\n", + "@trace\n", + "def generation(question: str, *context):\n", + " state: ProgramState = generation_sglang.run(question, *context)\n", + " while not state.stream_executor.is_finished:\n", + " time.sleep(1)\n", + " return state.stream_executor.variables[\"answer\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, we can tie it together and execute a sample query." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@trace\n", + "def rag_pipeline(question: str) -> str:\n", + " contexts = retrieval(question)\n", + " return generation(question, *contexts)\n", + "\n", + "\n", + "rag_pipeline(\n", + " \"When did the World Health Organization formally declare an end to the COVID-19 global health emergency?\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Debug Trace\n", + "\n", + "The output is unfortunately wrong! Using the traced pipeline, we can see that\n", + "\n", + "- the context is relevant to the question and contains the correct information\n", + "- but, the generation step is cut off as max tokens is set to 16\n", + "\n", + "When opening the generation step in the playground and rerunning the prompt with max. tokens set to 1000, the correct answer is produced.\n", + "\n", + "![RAG Trace](https://drive.google.com/uc?id=1QI243ogGjzbO01tUrR72g9rFoGzUJqVH)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Evaluating RAG Pipelines\n", + "\n", + "Before we apply above's fix, let's dive into evaluating RAG pipelines.\n", + "\n", + "RAG pipelines consist of a retrieval step to fetch relevant information and a generation step to generate a response to a users question. A RAG pipeline can fail at either step. E.g. the retrieval step can fail to find relevant information which makes generating the correct impossible. Another failure mode is that the generation step doesn't leverage the retrieved information correctly. We will apply the following evaluation metrics to understand different failure modes:\n", + "\n", + "- `context_relevancy`: measures how relevant the context is given the question\n", + "- `percent_target_supported_by_context`: measures how much of the target answer is supported by the context; this will give an upper ceiling of how well the generation step can perform\n", + "- `answer_context_faithfulness`: measures how much the generated answer utilizes the context\n", + "- `answer_matches_target`: measures how well the generated answer matches the target answer judged by a LLM and gives a sense of accuracy of our entire pipeline\n", + "\n", + "To use these evaluation metrics, we can import them from `parea.evals.rag` and `parea.evals.general` and apply them to a function by specifying in the `trace` decorator which evaluation metrics to use. The `@trace` decorator will automatically log the results of the evaluation metrics to the Parea dashboard.\n", + "\n", + "Applying them to the retrieval step:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from parea.evals.rag import (\n", + " context_query_relevancy_factory,\n", + " percent_target_supported_by_context_factory,\n", + ")\n", + "\n", + "\n", + "context_relevancy_eval = context_query_relevancy_factory()\n", + "percent_target_supported_by_context = percent_target_supported_by_context_factory()\n", + "\n", + "\n", + "@trace(eval_funcs=[context_relevancy_eval, percent_target_supported_by_context])\n", + "def retrieval(question: str) -> List[str]:\n", + " return collection.query(query_texts=[question], n_results=1)[\"documents\"][0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can apply `answer_context_faithfulness` and `answer_matches_target` to the generation step." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from parea.evals.general import answer_matches_target_llm_grader_factory\n", + "from parea.evals.rag import answer_context_faithfulness_statement_level_factory\n", + "\n", + "\n", + "answer_context_faithfulness = answer_context_faithfulness_statement_level_factory()\n", + "answer_matches_target_llm_grader = answer_matches_target_llm_grader_factory()\n", + "\n", + "\n", + "@function\n", + "def generation_sglang(s, question: str, *context: str):\n", + " context = \"\\n\".join(context)\n", + " s += user(\n", + " f\"Given this question:\\n{question}\\n\\nAnd this context:\\n{context}\\n\\nAnswer the question.\"\n", + " )\n", + " s += assistant(gen(\"answer\", max_tokens=1_000))\n", + "\n", + "\n", + "@trace(eval_funcs=[answer_context_faithfulness, answer_matches_target_llm_grader])\n", + "def generation(question: str, *context):\n", + " state: ProgramState = generation_sglang.run(question, *context)\n", + " while not state.stream_executor.is_finished:\n", + " time.sleep(1)\n", + " return state.stream_executor.variables[\"answer\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, we tie them together & execute the original sample query." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@trace\n", + "def rag_pipeline(question: str) -> str:\n", + " contexts = retrieval(question)\n", + " return generation(question, *contexts)\n", + "\n", + "\n", + "rag_pipeline(\n", + " \"When did the World Health Organization formally declare an end to the COVID-19 global health emergency?\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Great, the answer is correct! Can you spot the line where we fixed the output truncation issue?\n", + "\n", + "The evaluation scores appear in the bottom right of the logs (screenshot below). Note, that there is no score for `answer_matches_target_llm_grader` and `percent_target_supported_by_context` as these evals are automatically skipped if the target answer is not provided.\n", + "\n", + "![Fixed Max. Tokens](max-tokens-fixed-rag-trace.png)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Running an experiment\n", + "\n", + "Now we are (almost) ready to evaluate the performance of our RAG pipeline on the entire dataset. First, we will need to apply the `nest_asyncio` package to avoid issues with the Jupyter notebook event loop." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!pip install nest-asyncio\n", + "import nest_asyncio\n", + "\n", + "nest_asyncio.apply()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Running the actual experiment is straight-forward. For that we use `p.experiment` to initialize the experiment with a name, the data (list of key-value pairs fed into our entry function) and the entry function. We then call `run` on the experiment to execute it. Note, that `target` is a reserved key in the data dictionary and will be used as the target answer for evaluation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "e = p.experiment(\n", + " \"RAG\",\n", + " data=[\n", + " {\n", + " \"question\": qca[\"question\"],\n", + " \"target\": qca[\"answer\"],\n", + " }\n", + " for qca in question_context_answers\n", + " ],\n", + " func=rag_pipeline,\n", + ").run()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Analyzing the results\n", + "\n", + "When opening above experiment, we will see an overview of the experiment as shown below. The upper half shows a summary of the statistics on the left and charts to investigate the distribution and relationships of scores on the right. The lower half is a table with the individual traces which we can use to debug individual samples.\n", + "\n", + "When looking at the statistics, we can see that the accuracy of our RAG pipeline is 22% as measured by `answer_matches_target_llm_grader`. Though when checking the quality of our retrieval step (`context_query_relevancy`), we can see that our retrieval step is fetching relevant information in only 27% of all samples. As shown in the GIF, we investigate the relationship between the two and see the two scores have 95% agreement. This confirms that the retrieval step is a major bottleneck for our RAG pipeline. So, now it's your turn to improve the retrieval step!\n", + "\n", + "Note, above link isn't publicly accessible but the experiment can be accessed through [here](https://app.parea.ai/public-experiments/parea/rag_sglang/30f0244a-d56c-44ff-bdfb-8f47626304b6).\n", + "\n", + "![Experiment Results](https://drive.google.com/uc?id=1KMtJBU47nPB02Pvv3SPPTK7RnHRh5YdA)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/examples/frontend_language/usage/readme_examples.py b/examples/frontend_language/usage/readme_examples.py new file mode 100644 index 000000000..7269ef148 --- /dev/null +++ b/examples/frontend_language/usage/readme_examples.py @@ -0,0 +1,109 @@ +""" +Usage: +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +python readme_examples.py +""" + +import sglang as sgl + + +@sgl.function +def tool_use(s, question): + s += "To answer this question: " + question + ". " + s += ( + "I need to use a " + + sgl.gen("tool", choices=["calculator", "search engine"]) + + ". " + ) + + if s["tool"] == "calculator": + s += "The math expression is" + sgl.gen("expression") + elif s["tool"] == "search engine": + s += "The key word to search is" + sgl.gen("word") + + +@sgl.function +def tip_suggestion(s): + s += ( + "Here are two tips for staying healthy: " + "1. Balanced Diet. 2. Regular Exercise.\n\n" + ) + + forks = s.fork(2) + for i, f in enumerate(forks): + f += f"Now, expand tip {i+1} into a paragraph:\n" + f += sgl.gen(f"detailed_tip", max_tokens=256, stop="\n\n") + + s += "Tip 1:" + forks[0]["detailed_tip"] + "\n" + s += "Tip 2:" + forks[1]["detailed_tip"] + "\n" + s += "In summary" + sgl.gen("summary") + + +@sgl.function +def regular_expression_gen(s): + s += "Q: What is the IP address of the Google DNS servers?\n" + s += "A: " + sgl.gen( + "answer", + temperature=0, + regex=r"((25[0-5]|2[0-4]\d|[01]?\d\d?).){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)", + ) + + +@sgl.function +def text_qa(s, question): + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n") + + +def driver_tool_use(): + state = tool_use.run(question="What is the capital of the United States?") + print(state.text()) + print("\n") + + +def driver_tip_suggestion(): + state = tip_suggestion.run() + print(state.text()) + print("\n") + + +def driver_regex(): + state = regular_expression_gen.run() + print(state.text()) + print("\n") + + +def driver_batching(): + states = text_qa.run_batch( + [ + {"question": "What is the capital of the United Kingdom?"}, + {"question": "What is the capital of France?"}, + {"question": "What is the capital of Japan?"}, + ], + progress_bar=True, + ) + + for s in states: + print(s.text()) + print("\n") + + +def driver_stream(): + state = text_qa.run( + question="What is the capital of France?", temperature=0.1, stream=True + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print("\n") + + +if __name__ == "__main__": + # sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo-instruct")) + sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + + driver_tool_use() + driver_tip_suggestion() + driver_regex() + driver_batching() + driver_stream() diff --git a/examples/frontend_language/usage/sgl_gen_min_tokens.py b/examples/frontend_language/usage/sgl_gen_min_tokens.py new file mode 100644 index 000000000..a5088199b --- /dev/null +++ b/examples/frontend_language/usage/sgl_gen_min_tokens.py @@ -0,0 +1,35 @@ +""" +This example demonstrates how to use `min_tokens` to enforce sgl.gen to generate a longer sequence + +Usage: +python3 sgl_gen_min_tokens.py +""" + +import sglang as sgl + + +@sgl.function +def long_answer(s): + s += sgl.user("What is the capital of the United States?") + s += sgl.assistant(sgl.gen("answer", min_tokens=64, max_tokens=128)) + + +@sgl.function +def short_answer(s): + s += sgl.user("What is the capital of the United States?") + s += sgl.assistant(sgl.gen("answer")) + + +if __name__ == "__main__": + runtime = sgl.Runtime(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") + sgl.set_default_backend(runtime) + + state = long_answer.run() + print("=" * 20) + print("Longer Answer", state["answer"]) + + state = short_answer.run() + print("=" * 20) + print("Short Answer", state["answer"]) + + runtime.shutdown() diff --git a/examples/frontend_language/usage/streaming.py b/examples/frontend_language/usage/streaming.py new file mode 100644 index 000000000..506ee35c6 --- /dev/null +++ b/examples/frontend_language/usage/streaming.py @@ -0,0 +1,49 @@ +""" +Usage: +python3 streaming.py +""" + +import asyncio + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo")) + + +def stream_a_variable(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(var_name="answer_2"): + print(out, end="", flush=True) + print("\n") + + +async def async_stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + async for out in state.text_async_iter(var_name="answer_2"): + print(out, end="", flush=True) + print("\n") + + +if __name__ == "__main__": + stream_a_variable() + asyncio.run(async_stream()) diff --git a/examples/frontend_language/usage/triton/Dockerfile b/examples/frontend_language/usage/triton/Dockerfile new file mode 100644 index 000000000..e4741a1db --- /dev/null +++ b/examples/frontend_language/usage/triton/Dockerfile @@ -0,0 +1,10 @@ +FROM nvcr.io/nvidia/tritonserver:24.01-py3 + +WORKDIR /opt + +RUN git clone https://github.com/sgl-project/sglang.git + +WORKDIR /opt/sglang +RUN pip install --upgrade pip && \ + pip install -e "python[all]" && \ + pip install datasets diff --git a/examples/frontend_language/usage/triton/README.md b/examples/frontend_language/usage/triton/README.md new file mode 100644 index 000000000..b2e55961f --- /dev/null +++ b/examples/frontend_language/usage/triton/README.md @@ -0,0 +1,35 @@ +# sglang_triton + +Build the docker image: +``` +docker build -t sglang-triton . +``` + +Then do: +``` +docker run -ti --gpus=all --network=host --name sglang-triton -v ./models:/mnt/models sglang-triton +``` + +inside the docker container: +``` +cd sglang +python3 -m sglang.launch_server --model-path mistralai/Mistral-7B-Instruct-v0.2 --port 30000 --mem-fraction-static 0.9 +``` + +with another shell, inside the docker container: +``` +docker exec -ti sglang-triton /bin/bash +cd /mnt +tritonserver --model-repository=/mnt/models +``` + + +Send request to the server: +``` +curl -X POST http://localhost:8000/v2/models/character_generation/generate \ +-H "Content-Type: application/json" \ +-d '{ + "INPUT_TEXT": ["harry"] +}' + +``` diff --git a/examples/frontend_language/usage/triton/models/character_generation/config.pbtxt b/examples/frontend_language/usage/triton/models/character_generation/config.pbtxt new file mode 100644 index 000000000..7546f993a --- /dev/null +++ b/examples/frontend_language/usage/triton/models/character_generation/config.pbtxt @@ -0,0 +1,23 @@ +name: "character_generation" +backend: "python" +input [ + { + name: "INPUT_TEXT" + data_type: TYPE_STRING + dims: [ -1 ] + } +] +output [ + { + name: "OUTPUT_TEXT" + data_type: TYPE_STRING + dims: [ -1 ] + } +] +instance_group [ + { + count: 1 + kind: KIND_GPU + gpus: [ 0 ] + } +] diff --git a/examples/monitoring/README.md b/examples/monitoring/README.md new file mode 100644 index 000000000..3eef0b09b --- /dev/null +++ b/examples/monitoring/README.md @@ -0,0 +1,76 @@ +# SGLang Monitoring Setup + +This directory contains a ready-to-use monitoring setup for SGLang using Prometheus and Grafana. + +## Prerequisites + +- Docker and Docker Compose installed +- SGLang server running with metrics enabled + +## Usage + +1. Start your SGLang server with metrics enabled: + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --port 30000 --enable-metrics +``` + +By default, the metrics server will run on `127.0.0.1:30000`. + +2. Start the monitoring stack: + +```bash +cd examples/monitoring +docker compose up +``` + +3. Access the monitoring interfaces: + - Grafana: [http://localhost:3000](http://localhost:3000) + - Prometheus: [http://localhost:9090](http://localhost:9090) + +Default Grafana login credentials: +- Username: `admin` +- Password: `admin` + +You'll be prompted to change the password on first login. + +4. The SGLang dashboard will be automatically available in the "SGLang Monitoring" folder. + +## Troubleshooting + +### Port Conflicts +If you see errors like "port is already allocated": + +1. Check if you already have Prometheus or Grafana running: + ```bash + docker ps | grep -E 'prometheus|grafana' + ``` + +2. Stop any conflicting containers: + ```bash + docker stop + ``` + +3. Ensure no other services are using ports 9090 and 3000: + ```bash + lsof -i :9090 + lsof -i :3000 + ``` + +### Connection Issues +If Grafana cannot connect to Prometheus: +1. Check that both services are running +2. Verify the datasource configuration in Grafana +3. Check that your SGLang server is properly exposing metrics + +## Configuration + +- Prometheus configuration: `prometheus.yaml` +- Docker Compose configuration: `docker-compose.yaml` +- Grafana datasource: `grafana/datasources/datasource.yaml` +- Grafana dashboard configuration: `grafana/dashboards/config/dashboard.yaml` +- SGLang dashboard JSON: `grafana/dashboards/json/sglang-dashboard.json` + +## Customization + +You can customize the monitoring setup by modifying the configuration files as needed. diff --git a/examples/monitoring/docker-compose.yaml b/examples/monitoring/docker-compose.yaml new file mode 100644 index 000000000..ce6457fa3 --- /dev/null +++ b/examples/monitoring/docker-compose.yaml @@ -0,0 +1,28 @@ +version: '3' +services: + prometheus: + image: prom/prometheus:latest + container_name: prometheus + network_mode: host + volumes: + - ./prometheus.yaml:/etc/prometheus/prometheus.yml + command: + - '--config.file=/etc/prometheus/prometheus.yml' + - '--storage.tsdb.path=/prometheus' + + grafana: + image: grafana/grafana:latest + container_name: grafana + network_mode: host + volumes: + - ./grafana/datasources:/etc/grafana/provisioning/datasources + - ./grafana/dashboards/config:/etc/grafana/provisioning/dashboards + - ./grafana/dashboards/json:/var/lib/grafana/dashboards + environment: + - GF_AUTH_ANONYMOUS_ENABLED=true + - GF_AUTH_ANONYMOUS_ORG_ROLE=Viewer + - GF_AUTH_BASIC_ENABLED=false + - GF_USERS_ALLOW_SIGN_UP=false + - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/sglang-dashboard.json + depends_on: + - prometheus diff --git a/examples/monitoring/grafana/dashboards/config/dashboard.yaml b/examples/monitoring/grafana/dashboards/config/dashboard.yaml new file mode 100644 index 000000000..6c17a6c63 --- /dev/null +++ b/examples/monitoring/grafana/dashboards/config/dashboard.yaml @@ -0,0 +1,11 @@ +apiVersion: 1 +providers: + - name: 'SGLang' + orgId: 1 + folder: 'SGLang Monitoring' + type: file + disableDeletion: false + updateIntervalSeconds: 10 + allowUiUpdates: false + options: + path: /var/lib/grafana/dashboards diff --git a/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json b/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json new file mode 100644 index 000000000..ebd663e41 --- /dev/null +++ b/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json @@ -0,0 +1,984 @@ +{ + "annotations": { + "list": [ + { + "builtIn": 1, + "datasource": { + "type": "grafana", + "uid": "-- Grafana --" + }, + "enable": true, + "hide": true, + "iconColor": "rgba(0, 211, 255, 1)", + "name": "Annotations & Alerts", + "type": "dashboard" + } + ] + }, + "editable": true, + "fiscalYearStartMonth": 0, + "graphTooltip": 0, + "id": 8, + "links": [], + "panels": [ + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 0 + }, + "id": 14, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.99, sum by (le) (rate(sglang:e2e_request_latency_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P99", + "range": true, + "refId": "A", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.9, sum by (le) (rate(sglang:e2e_request_latency_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P90", + "range": true, + "refId": "B", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.5, sum by (le) (rate(sglang:e2e_request_latency_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P50", + "range": true, + "refId": "C", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "avg(rate(sglang:e2e_request_latency_seconds_sum[$__rate_interval]) / rate(sglang:e2e_request_latency_seconds_count[$__rate_interval]))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "Avg", + "range": true, + "refId": "D", + "useBackend": false + } + ], + "title": "End-to-End Request Latency", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "custom": { + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "scaleDistribution": { + "type": "linear" + } + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 0 + }, + "id": 17, + "maxDataPoints": 30, + "options": { + "calculate": false, + "calculation": { + "yBuckets": { + "scale": { + "type": "linear" + } + } + }, + "cellGap": 1, + "cellValues": {}, + "color": { + "exponent": 0.5, + "fill": "dark-orange", + "mode": "scheme", + "reverse": false, + "scale": "exponential", + "scheme": "Spectral", + "steps": 64 + }, + "exemplars": { + "color": "rgba(255,0,255,0.7)" + }, + "filterValues": { + "le": 1e-9 + }, + "legend": { + "show": true + }, + "rowsFrame": { + "layout": "auto" + }, + "tooltip": { + "mode": "single", + "showColorScale": true, + "yHistogram": false + }, + "yAxis": { + "axisPlacement": "left", + "reverse": false, + "unit": "secs" + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "sum(increase(sglang:e2e_request_latency_seconds_bucket{model_name=~\"$model_name\"}[$__rate_interval])) by (le)\r\n", + "format": "heatmap", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "{{le}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "End-to-End Request Latency(s) Heatmap", + "type": "heatmap" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 8 + }, + "id": 20, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.99, sum by (le) (rate(sglang:time_to_first_token_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P99", + "range": true, + "refId": "A", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.9, sum by (le) (rate(sglang:time_to_first_token_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P90", + "range": true, + "refId": "B", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.5, sum by (le) (rate(sglang:time_to_first_token_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P50", + "range": true, + "refId": "C", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "avg(rate(sglang:time_to_first_token_seconds_sum[$__rate_interval]) / rate(sglang:time_to_first_token_seconds_count[$__rate_interval]))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "Avg", + "range": true, + "refId": "D", + "useBackend": false + } + ], + "title": "Time-To-First-Token Latency", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "custom": { + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "scaleDistribution": { + "type": "linear" + } + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 8 + }, + "id": 19, + "maxDataPoints": 30, + "options": { + "calculate": false, + "calculation": { + "xBuckets": { + "value": "" + }, + "yBuckets": { + "mode": "size", + "scale": { + "type": "linear" + }, + "value": "" + } + }, + "cellGap": 1, + "color": { + "exponent": 0.5, + "fill": "dark-orange", + "mode": "scheme", + "reverse": false, + "scale": "exponential", + "scheme": "Spectral", + "steps": 64 + }, + "exemplars": { + "color": "rgba(255,0,255,0.7)" + }, + "filterValues": { + "le": 1e-9 + }, + "legend": { + "show": true + }, + "rowsFrame": { + "layout": "auto" + }, + "tooltip": { + "mode": "single", + "showColorScale": true, + "yHistogram": false + }, + "yAxis": { + "axisPlacement": "left", + "reverse": false + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "builder", + "exemplar": false, + "expr": "sum by(le) (increase(sglang:time_to_first_token_seconds_bucket{model_name=~\"$model_name\"}[$__rate_interval]))", + "format": "heatmap", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "interval": "", + "legendFormat": "{{le}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "Time-To-First-Token Seconds Heatmap", + "type": "heatmap" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 16 + }, + "id": 7, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "sglang:num_running_reqs", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "interval": "", + "legendFormat": "{{instance}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "Num Running Requests", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 16 + }, + "id": 18, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "editorMode": "code", + "expr": "sglang:gen_throughput", + "instant": false, + "legendFormat": "{{instance}}", + "range": true, + "refId": "A" + } + ], + "title": "Token Generation Throughput (Tokens / S)", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 24 + }, + "id": 11, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "sglang:cache_hit_rate", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "{{instance}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "Cache Hit Rate", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 24 + }, + "id": 8, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "sglang:num_queue_reqs", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "{{instance}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "Number Queued Requests", + "type": "timeseries" + } + ], + "preload": false, + "refresh": "5s", + "schemaVersion": 41, + "tags": [], + "templating": { + "list": [ + { + "current": { + "text": "127.0.0.1:30000", + "value": "127.0.0.1:30000" + }, + "datasource": { + "type": "prometheus" + }, + "definition": "label_values(instance)", + "includeAll": false, + "label": "instance", + "name": "instance", + "options": [], + "query": { + "qryType": 1, + "query": "label_values(instance)", + "refId": "PrometheusVariableQueryEditor-VariableQuery" + }, + "refresh": 1, + "regex": "", + "type": "query" + }, + { + "current": { + "text": "meta-llama/Llama-3.1-8B-Instruct", + "value": "meta-llama/Llama-3.1-8B-Instruct" + }, + "datasource": { + "type": "prometheus" + }, + "definition": "label_values(model_name)", + "includeAll": false, + "label": "model name", + "name": "model_name", + "options": [], + "query": { + "qryType": 1, + "query": "label_values(model_name)", + "refId": "PrometheusVariableQueryEditor-VariableQuery" + }, + "refresh": 1, + "regex": "", + "type": "query" + } + ] + }, + "time": { + "from": "now-30m", + "to": "now" + }, + "timepicker": {}, + "timezone": "browser", + "title": "SGLang Dashboard", + "uid": "sglang-dashboard", + "version": 11 +} diff --git a/examples/monitoring/grafana/datasources/datasource.yaml b/examples/monitoring/grafana/datasources/datasource.yaml new file mode 100644 index 000000000..1ab0e4a5f --- /dev/null +++ b/examples/monitoring/grafana/datasources/datasource.yaml @@ -0,0 +1,8 @@ +apiVersion: 1 +datasources: + - name: Prometheus + type: prometheus + access: proxy + url: http://localhost:9090 + isDefault: true + editable: false diff --git a/examples/monitoring/prometheus.yaml b/examples/monitoring/prometheus.yaml new file mode 100644 index 000000000..ba16ac3bd --- /dev/null +++ b/examples/monitoring/prometheus.yaml @@ -0,0 +1,10 @@ +# prometheus.yaml +global: + scrape_interval: 5s + evaluation_interval: 30s + +scrape_configs: + - job_name: sglang + static_configs: + - targets: + - '127.0.0.1:30000' diff --git a/examples/profiler/nsys_profile_tools/README.md b/examples/profiler/nsys_profile_tools/README.md new file mode 100644 index 000000000..687200e05 --- /dev/null +++ b/examples/profiler/nsys_profile_tools/README.md @@ -0,0 +1,176 @@ +# gputrc2graph.py + +This script processes NVIDIA Nsight Systems (`nsys`) GPU trace files +(`.nsys-rep`) with -t cuda tracing enabled, and generates kernel-level +summaries and visualizations of GPU and non-GPU time. It is useful for +profiling and analyzing nsys profile output. + +## Usage + +### Command-line Arguments + +- `--in_file` + **(required)** + List of input files and their metadata. Each entry should be in the format: + `,,,` + - `nsys-rep`: Path to the `.nsys-rep` file. + - `engine`: Engine name (e.g., `sglang`). + - `model`: Model name (e.g., `llama`, `gpt-oss`, `ds`). + - `elapsed_nonprofiled_sec`: Wall-clock runtime (in seconds) without + profiling. Specify `0` to use the elapsed time from the nsys-rep file + (this may inflate non-GPU time if actual runtime without profiling is + less). Multiple entries can be provided, separated by spaces. + +- `--out_dir` + Output directory for the generated CSV and HTML files. + If not specified, results are saved in the current directory. + +- `--title` + Title for the HTML chart/visualization. + +- `--nsys_cmd` + Path to the `nsys` command. + Default: `nsys` (assumes it is in your PATH). + Use this if `nsys` is not in your system PATH. + +## Notes + +- Make sure you have pandas installed. Any version is fine. +- Make sure [nsys](https://developer.nvidia.com/nsight-systems/get-started) is +installed, and specify the path to the `nsys` command with `--nsys_cmd` if it + is not in your PATH. The nsys version must be >= the nsys profile version that + was used to collect the traces when profiling the server, so that nsys can + process the nsys-rep that was generated. + +- For more details on available engines and models, see the help string in + the script or run: + +```bash +python3 gputrc2graph.py --help +``` + +## Example 1: analyze a single profile + +To analyze the GPU cycles of for example, a llama-3.1-8B model with sglang: + +1. Run the following command to collect nsys profile, for sglang server config. + + ```bash + nsys profile -t cuda -o nsys_res -f true --trace-fork-before-exec=true \ + --cuda-graph-trace=node --delay --duration \ + python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B ... + ``` + + where: + + - DELAY: how many seconds to delay nsys from collecting profiles, needed so + that profiles aren't captured till sglang server has come up and load + generation starts. + - DURATION: how many seconds for nsys profile to run before generating the + profile. This should be > the duration of the run. +2. After the server starts, run the client load generation command. Once the +test completes, after DURATION amount of time, nsys profile will generate an +nsys_res.nsys-rep file and shut down the server. + +3. Run step #1 again, this time starting up the server without collecting the +profile. + +4. Run step #2 again, and record the total time to complete the test in +seconds. This value will be used by the script to calculate the + CPU(non-GPU) seconds for the analysis. + +5. Say the run elapsed time from step #4 is 132 seconds. Run script to + analyze: + + ```bash + python3 gputrc2graph.py \ + --in_file run1.nsys-rep,sglang,llama,132 + ``` + +The command will produce 2 files for analysis: + +- result.html: this categorizes kernel names into different categories in a + stacked bar chart. +- result.csv: shows how the kernel names are mapped to the different + categories. + +### HTML visualization with result.html + +The html file shows the number of elapsed seconds due to different GPU +Substages or categories, which consist of attention kernels as the biggest +category, at 63 seconds, followed by "gemm" kernels. This lets the user +prioritize the kernels to focus on for performance optimizations. + +There's also an appended data table underneath the bar chart for copying out to + other post-processing tools. + +### Kernel to category mapping with result.csv + +Suppose the user would like to focus on improving triton kernels. It's not the +biggest consumer of cycles at .01 sec but perhaps it hasn't been optimized. +The next step is to use the result.csv to dive into what the kernels are which +compose the triton kernel GPU cycles. + +## Example 2: analyze multiple profiles + +Suppose the user has multiple nsys trace files, captured for different models, +say llama and gpt-oss in this case, and wish to compare their GPU/non-GPU +time, something like the following command can be used. + +```bash +python3 gputrc2graph.py \ +--in_file run1.nsys-rep,sglang,llama,100 run2.nsys-rep,sglang,gpt-oss,102 \ +--out_dir results +``` + +The analysis process is similar to example 1 but now there will be multiple +stack bar charts that can be compared. The categories for the different +kernels will remain the same, so that it's easy to compare the GPU cycles for +the same categories. + +Once a category is shown to have more cycles for one configuration than +another, the next step would be to use the csv file to see what kernels are +mapped into that category, and which kernels are taking the largest amount of +time which would cause a difference for the overall category. + +## Example 3: add new classification for a new model + +To create a new engine DEF with model ABC, just add another json file in the same directory as +gputrc2graph.py with the same format as the other json files. The script will automatically pick up all the json files in the same directory as engine/model specifications. + +Then, for this new model, suppose there are 4 kernels to be classified into +"gemm" and "attn", where the gemm kernels have names with "*H*" or "*I*" in +them, and attn kernels have names with "*J*" or "*K*" in them, just add another + .json file in the same directory as gputrc2graph.py with the same format as + the other json files, like the following: + +```json +{ + "DEF": { + "ABC": { + "H|I": "gemm", + "J|K": "attn", + "CUDA mem": "non-gpu-H_D_memops", + ".*": "misc" + } + } +} +``` + +Each entry in the dictionary consists of: + +- key: a regex used to classify the kernels +- value: the category to classify the kernels into. + +The last 2 entries are common for all engine/models, consisting of CUDA memory +operations and a 'misc' for anything that's leftover and can't be classified. + +When invoking gputrc2graph.py, specify a trace file with this new model/engine +like the following: + +```bash +--in_file new.nsys-rep,DEF,ABC, +``` + +If the engine_DEF.json file already exists, just add the model as a new node in + the existing engine file, after the other models. diff --git a/examples/profiler/nsys_profile_tools/gputrc2graph.py b/examples/profiler/nsys_profile_tools/gputrc2graph.py new file mode 100755 index 000000000..f17bd1857 --- /dev/null +++ b/examples/profiler/nsys_profile_tools/gputrc2graph.py @@ -0,0 +1,344 @@ +""" + This generates gpu kernel analysis output from nsys rep. Will call nsys + stats -r cuda_gpu_kern_trace, get non-overlapped gpu cycles, then generate + csv and html output for analysis +""" + +import argparse +import logging +import os + +import regex as re + +logger = logging.getLogger(__name__) + + +# helper data class for annotating kernels +def load_engine_model(): + """returns engine_model built from all json files in the current dir""" + import glob + import json + + engine_model = {} + + json_files = glob.glob(os.path.join(os.path.dirname(__file__) or ".", "*.json")) + for fname in json_files: + with open(fname, encoding="utf-8") as f: + engine_model.update(json.load(f)) + return engine_model + + +class GPUTrace2Graph: + """ + Parses output of nsys report, generates csv and bar chart output + """ + + def __init__(self): + import pandas as pd # avoid importing till needed + + self.pd = pd + self.pd.options.mode.copy_on_write = True + + # helper functions for generating trace->summary csvs + def gen_nonoverlapped_sum_from_gputrace(self, in_file, out_file): + logger.info("loading %s", in_file) + df = self.pd.read_csv( + in_file, usecols=["Start (ns)", "Duration (ns)", "Device", "Strm", "Name"] + ) + df["End (ns)"] = df["Start (ns)"] + df["Duration (ns)"] + df = self.sum_non_overlapping_intervals(df) + # get ready to print table with elapsed times per kernel + df["Instances"] = 1 + df_sum = df.groupby("Name", as_index=False).agg( + {"Elapsed Time (ns)": "sum", "Duration (ns)": "sum", "Instances": "size"} + ) + + # generate csv + df_sum["Total Time (sec)"] = df_sum["Duration (ns)"] / 1e9 + df_sum["Elapsed Time (sec)"] = df_sum["Elapsed Time (ns)"] / 1e9 + df_sum = df_sum.sort_values(by="Elapsed Time (sec)", ascending=False) + df_sum[["Elapsed Time (sec)", "Total Time (sec)", "Instances", "Name"]].to_csv( + out_file, index=False + ) + + def sum_non_overlapping_intervals(self, df): + """ + returns new sorted df with Elapsed Time (ns) column using + vectorized operations + """ + logger.info("sorting %s trace records by start time", str(df.shape)) + + # Sort by start time and reset index + df = df.sort_values(by="Start (ns)").reset_index(drop=True) + + # Initialize elapsed time as duration + df["Elapsed Time (ns)"] = df["Duration (ns)"] + + # Get numpy arrays for faster operations + starts = df["Start (ns)"].values + ends = df["End (ns)"].values + + # Keep track of current interval end + current_end = ends[0] + display_units = max(1, int(len(df) / 100)) + # Update current_end for overlapping intervals + for i in range(1, len(df)): + if i % display_units == 0: + print(f"processing trace: {int(i/len(df) * 100)} %", end="\r") + if starts[i] <= current_end: + if ends[i] > current_end: + # Partial overlap + df.iloc[i, df.columns.get_loc("Elapsed Time (ns)")] = ( + ends[i] - current_end + ) + current_end = ends[i] + else: + # Complete overlap + df.iloc[i, df.columns.get_loc("Elapsed Time (ns)")] = 0 + else: + # No overlap + current_end = ends[i] + + return df + + # functions for generating html files + def make_html(self, df, output_dir, title): + """make html graph from df""" + import plotly.express as px + + if df.empty: + return + output_name = os.path.join(output_dir, "result") + if not title: + title = "Model_Engine" + x = "Model_Engine" + y = "Elapsed Time (sec)" + color = "Category" + """ generate kernel mapping table """ + # Sort Model_Engine categories by last field after underscore + df["Model_Engine"] = self.pd.Categorical( + df["Model_Engine"], + sorted(df["Model_Engine"].unique(), key=lambda x: x.split("_")[-1]), + ) + df[["Model_Engine", color, "Instances", "Name", y]].sort_values( + by=color + ).to_csv(f"{output_name}.csv", index=False) + graph = px.histogram( + df.round(2), + x=x, + y=y, + title=(f"{y} for {title}"), + color=color, + text_auto=True, + ) + # wrap x axis labels + graph.update_xaxes(automargin=True) + graph.write_html(f"{output_name}.html") + """ + Generate data table with columns per Model_Engine into result.html + """ + pivot_df = df.pivot_table( + values="Elapsed Time (sec)", + index="Category", + columns="Model_Engine", + aggfunc="sum", + observed=False, + ).round(2) + # Add sum row at bottom + pivot_df.loc["total_elapsed_sec"] = pivot_df.sum() + pivot_df.fillna("").to_html("temp.html") + with ( + open(f"{output_name}.html", "a", encoding="utf-8") as outfile, + open("temp.html", encoding="utf-8") as infile, + ): + outfile.write(infile.read()) + os.remove("temp.html") + + print( + f"Finished generating: \n" + f" {output_name}.html for stack bar chart \n" + f" {output_name}.csv for Kernel-Category mapping" + ) + + def anno_gpu_kernname(self, df, mapping): + """add "Category" column""" + + def anno_gpu_kernname_helper(name): + for kern_name, val in mapping.items(): + if re.search(kern_name, name): + return val + + df["Category"] = df["Name"].apply(anno_gpu_kernname_helper) + + def make_nongpu_row(self, df, nongpu_sec): + """this will append non-gpu time entry at end of df""" + nongpu_row = self.pd.DataFrame([df.iloc[-1]]) + nongpu_row["Category"] = nongpu_row["Name"] = "CPU(non-GPU)" + nongpu_row["Instances"] = 1 + nongpu_row["Elapsed Time (sec)"] = nongpu_sec + return nongpu_row + + def is_valid_file(self, base_file): + """asserts if base_file is non-existent or is empty""" + assert ( + os.path.isfile(base_file) and os.path.getsize(base_file) > 0 + ), f"{base_file} doesn't exist or is empty" + + def should_gen_file(self, new_file, base_file): + """figure out if new file should be generated from base_file""" + self.is_valid_file(base_file) + if ( + os.path.exists(new_file) + and (os.path.getmtime(new_file) > os.path.getmtime(base_file)) + and (os.path.getsize(base_file) > 0) + ): + logger.info("reusing %s", new_file) + return False + else: + logger.info("generating %s", new_file) + return True + + def gen_sum_file(self, file, nsys_cmd): + """ + generates sum file from nsys trace with times per kernel and + returns the name of the sum file + """ + import subprocess + + file_dir = os.path.dirname(file) + file_name = os.path.basename(file) + + if not file_dir: + file_dir = "." + # Walk through trace and get the total non-overlapped time + nsys_stats_file = os.path.join(file_dir, f"{file_name}_cuda_gpu_trace.csv") + sum_file = os.path.join(file_dir, f"{file_name}_cuda_gpu_kernel_tracesum.csv") + if self.should_gen_file(nsys_stats_file, file): + cmd = [ + nsys_cmd, + "stats", + "-r", + "cuda_gpu_trace", + file, + "-o", + f"{file_dir}/{file_name}", + ] + cmd_str = " ".join(cmd) + logger.info("+ %s", cmd_str) + # estimate time based on calibrated 240M/min + file_size_mb = os.path.getsize(file) / 1e6 + logger.info( + "nsys stats for %.2f MB file expected to take %.2f min", + file_size_mb, + file_size_mb / 240, + ) + try: + subprocess.run(cmd, check=True) + except (FileNotFoundError, subprocess.CalledProcessError) as e: + logger.error( + "'%s' failed: %s. Use --nsys_cmd to specify nsys path", cmd_str, e + ) + exit(1) + logger.info("generating non-overalapped sum %s", sum_file) + self.gen_nonoverlapped_sum_from_gputrace(nsys_stats_file, sum_file) + self.is_valid_file(sum_file) + logger.info("Finished generating %s", sum_file) + return sum_file + + def gen_graph(self, in_file, out_dir, title, nsys_cmd, engine_model): + """generates graph and csv file from in_file into out_dir""" + # Initialize an empty DataFrame to store combined data + combined_df = self.pd.DataFrame() + for idx, (file, engine, model, total_sec) in enumerate(in_file): + file_dir = os.path.dirname(file) + file_name = os.path.basename(file) + if not file_dir: + file_dir = "." + sum_file = self.gen_sum_file(file, nsys_cmd) + # read kernel summary file + df = self.pd.read_csv(sum_file) + # annotate kernel to their categories + assert engine_model.get(engine), f"engine {engine} unknown" + assert engine_model[engine].get(model), f"model {model} unknown" + # remove nsys-rep from file_name for shorter x-label + file_name = file_name.replace(".nsys-rep", "") + df["Model_Engine"] = f"{model}_{engine}_{file_name}_{idx}" + self.anno_gpu_kernname(df, engine_model[engine][model]) + # patch in non-gpu time + gpu_sec = round(df["Elapsed Time (sec)"].sum(), 1) + total_sec = round(float(total_sec), 1) + if total_sec < gpu_sec: + logger.warning( + "Elapsed sec %.2f < GPU sec %.2f resetting Elapsed sec ", + total_sec, + gpu_sec, + ) + total_sec = gpu_sec + nongpu_row = self.make_nongpu_row(df, total_sec - gpu_sec) + df = self.pd.concat([df, nongpu_row], ignore_index=True) + combined_df = self.pd.concat([combined_df, df], ignore_index=True) + if out_dir is None: + out_dir = "." + else: + os.makedirs(out_dir, exist_ok=True) + # generate html file + self.make_html(combined_df, out_dir, title) + + +def parse_tuple(s): + return tuple(s.split(",")) + + +def main(): + logging.basicConfig( + format=("%(asctime)s - %(levelname)s - %(message)s"), level=logging.INFO + ) + parser = argparse.ArgumentParser( + description=( + "Process nsys rep and generate kernel non-overlapped cycles. \n" + "Example:\n" + "gputrc2graph.py --in_file d1.nsys-rep,sglang,llama,100 \n" + "d2.nsys-rep,sglang,gpt-oss,102 " + '--out_dir results/ --title "Model=gpt-oss SGLANG chart"' + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + # load supported engine_model + engine_model_supported = load_engine_model() + # Get a string representation of supported engine/model combinations + engine_model_supported_str = ", ".join( + f"{engine}:[{', '.join(models.keys())}]" + for engine, models in engine_model_supported.items() + ) + parser.add_argument( + "--in_file", + type=parse_tuple, + nargs="+", + help=( + "list of (nsys-rep, engine, model, elapsed_nonprofiled_sec) " + "separated by space. Elapsed_nonprofiled_sec is runtime without " + "profiling used to calculate non-gpu time. Specify 0 to use " + "elapsed time from nsys-rep but that might inflate non-gpu time. " + f"Available engine:[model] are: {engine_model_supported_str} " + f"Example: --infile d1.nsys-rep,sglan,llama,100 " + "d2.nsys-rep,sglang,gpt-oss,102" + ), + required=True, + ) + parser.add_argument("--out_dir", help=("output dir for result.csv/html")) + parser.add_argument("--title", help=("title for html chart")) + parser.add_argument( + "--nsys_cmd", + help=("nsys cmd, e.g. /usr/bin/nsys, Default: nsys"), + default="nsys", + ) + args = parser.parse_args() + gputrace = GPUTrace2Graph() + gputrace.gen_graph( + args.in_file, args.out_dir, args.title, args.nsys_cmd, engine_model_supported + ) + + +if __name__ == "__main__": + main() diff --git a/examples/profiler/nsys_profile_tools/sglang_engine_model.json b/examples/profiler/nsys_profile_tools/sglang_engine_model.json new file mode 100644 index 000000000..253cc762b --- /dev/null +++ b/examples/profiler/nsys_profile_tools/sglang_engine_model.json @@ -0,0 +1,61 @@ +{ + "sglang": { + "llama": { + "gemm|nvjet": "gemm", + "fused_moe_kernel|GroupProblemShape|group_gemm_starts|bmm_|GemmUniversal": "moe_gemm", + "moe|sigmoid": "moe", + "CatArrayBatched|prepare_inputs": "prepare_next", + "ncclDevKernel|cross_device_reduce": "nccl_and_custom_ar", + "_norm_|Norm": "norm", + "topk": "topk", + "act_and_mul_": "activation", + "Rotary": "rope", + "SoftMax": "softmax", + "flash|fmha": "attn", + "elementwise": "elementwise", + "fp8_quant|cvt_|quantize": "quantize", + "reduce_kernel": "reduce", + "triton": "triton_kernel", + "CUDA mem": "non-gpu-H_D_memops", + ".*": "misc" + }, + "ds": { + "block_fp8_matmul": "block_fp8_gemm", + "gemm|matmul|nvjet": "gemm", + "fused_moe_kernel": "moe_gemm", + "moe|expert|sigmoid": "moe", + "CatArrayBatched|write_req_to": "prepare_next", + "ncclDevKernel|cross_device_reduce|all_gather": "nccl_and_custom_ar", + "Norm": "norm", + "topk": "topk", + "activation|act_and_mul": "activation", + "compute_position_kernel": "rope", + "elementwise": "elementwise", + "fp8_quant|quant_fp8|quantize": "quantize", + "SoftMax": "softmax", + "reduce": "reduce", + "_fwd_|create_flash|::mla::|KVCache": "attn", + "CUDA mem": "non-gpu-H_D_memops", + ".*": "misc" + }, + "gpt-oss": { + "gemm|nvjet": "gemm", + "fused_moe_kernel|_group_gemm|GroupProblemShape|GemmUniversal|bmm_|matmul_ogs_|_topk_forward|_combined_routing|_sum_bitmatrix_rows|_compute_writeback_idx": "moe_gemm", + "moe|sigmoid": "moe", + "CatArrayBatched|prepare_inputs": "prepare_next", + "_norm_|Norm": "norm", + "ncclDevKernel|cross_device_reduce|allreduce": "nccl_and_custom_ar", + "topk|TopK": "topk", + "act_and_mul_": "activation", + "Rotary": "rope", + "SoftMax": "softmax", + "flash|fmha": "attn", + "elementwise": "elementwise", + "fp8_quant|cvt_|quantize": "quantize", + "reduce_kernel": "reduce", + "triton": "triton_kernel", + "CUDA mem": "non-gpu-H_D_memops", + ".*": "misc" + } + } +} diff --git a/examples/runtime/README.md b/examples/runtime/README.md new file mode 100644 index 000000000..18414452f --- /dev/null +++ b/examples/runtime/README.md @@ -0,0 +1,45 @@ +# Runtime examples + +The below examples will mostly need you to start a server in a separate terminal before you can execute them. Please see in the code for detailed instruction. + +## Native API + +* `lora.py`: An example how to use LoRA adapters. +* `multimodal_embedding.py`: An example how perform [multi modal embedding](Alibaba-NLP/gme-Qwen2-VL-2B-Instruct). +* `openai_batch_chat.py`: An example how to process batch requests for chat completions. +* `openai_batch_complete.py`: An example how to process batch requests for text completions. +* **`openai_chat_with_response_prefill.py`**: + An example that demonstrates how to [prefill a response](https://eugeneyan.com/writing/prompting/#prefill-claudes-responses) using the OpenAI API by enabling the `continue_final_message` parameter. + When enabled, the final (partial) assistant message is removed and its content is used as a prefill so that the model continues that message rather than starting a new turn. See [Anthropic's prefill example](https://docs.anthropic.com/en/docs/build-with-claude/prompt-engineering/prefill-claudes-response#example-structured-data-extraction-with-prefilling) for more context. +* `reward_model.py`: An example how to extract scores from a reward model. +* `vertex_predict.py`: An example how to deploy a model to [Vertex AI](https://cloud.google.com/vertex-ai?hl=en). + +## Engine + +The `engine` folder contains that examples that show how to use [Offline Engine API](https://docs.sglang.ai/backend/offline_engine_api.html#Offline-Engine-API) for common workflows. + +* `custom_server.py`: An example how to deploy a custom server. +* `embedding.py`: An example how to extract embeddings. +* `launch_engine.py`: An example how to launch the Engine. +* `offline_batch_inference_eagle.py`: An example how to perform speculative decoding using [EAGLE](https://docs.sglang.ai/backend/speculative_decoding.html). +* `offline_batch_inference_torchrun.py`: An example how to perform inference using [torchrun](https://pytorch.org/docs/stable/elastic/run.html). +* `offline_batch_inference_vlm.py`: An example how to use VLMs with the engine. +* `offline_batch_inference.py`: An example how to use the engine to perform inference on a batch of examples. + +## Hidden States + +The `hidden_states` folder contains examples on how to extract hidden states using SGLang. Please note that this might degrade throughput due to cuda graph rebuilding. + +* `hidden_states_engine.py`: An example how to extract hidden states using the Engine API. +* `hidden_states_server.py`: An example how to extract hidden states using the Server API. + +## Multimodal + +SGLang supports multimodal inputs for various model architectures. The `multimodal` folder contains examples showing how to use urls, files or encoded data to make requests to multimodal models. Examples include querying the [Llava-OneVision](multimodal/llava_onevision_server.py) model (image, multi-image, video), Llava-backed [Qwen-Llava](multimodal/qwen_llava_server.py) and [Llama3-Llava](multimodal/llama3_llava_server.py) models (image, multi-image), and Mistral AI's [Pixtral](multimodal/pixtral_server.py) (image, multi-image). + + +## Token In, Token Out + +The folder `token_in_token_out` shows how to perform inference, where we provide tokens and get tokens as response. + +* `token_in_token_out_{llm|vlm}_{engine|server}.py`: Shows how to perform token in, token out workflow for llm/vlm using either the engine or native API. diff --git a/examples/runtime/engine/custom_server.py b/examples/runtime/engine/custom_server.py new file mode 100644 index 000000000..b190a463e --- /dev/null +++ b/examples/runtime/engine/custom_server.py @@ -0,0 +1,53 @@ +from sanic import Sanic, text +from sanic.response import json + +import sglang as sgl + +engine = None + +# Create an instance of the Sanic app +app = Sanic("sanic-server") + + +# Define an asynchronous route handler +@app.route("/generate", methods=["POST"]) +async def generate(request): + prompt = request.json.get("prompt") + if not prompt: + return json({"error": "Prompt is required"}, status=400) + + # async_generate returns a dict + result = await engine.async_generate(prompt) + + return text(result["text"]) + + +@app.route("/generate_stream", methods=["POST"]) +async def generate_stream(request): + prompt = request.json.get("prompt") + + if not prompt: + return json({"error": "Prompt is required"}, status=400) + + # async_generate returns a dict + result = await engine.async_generate(prompt, stream=True) + + # https://sanic.dev/en/guide/advanced/streaming.md#streaming + # init the response + response = await request.respond() + + # result is an async generator + async for chunk in result: + await response.send(chunk["text"]) + + await response.eof() + + +def run_server(): + global engine + engine = sgl.Engine(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") + app.run(host="0.0.0.0", port=8000, single_process=True) + + +if __name__ == "__main__": + run_server() diff --git a/examples/runtime/engine/embedding.py b/examples/runtime/engine/embedding.py new file mode 100644 index 000000000..b927a188b --- /dev/null +++ b/examples/runtime/engine/embedding.py @@ -0,0 +1,27 @@ +import sglang as sgl + + +def main(): + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + # Create an LLM. + llm = sgl.Engine( + model_path="Alibaba-NLP/gte-Qwen2-1.5B-instruct", is_embedding=True + ) + + outputs = llm.encode(prompts) + # Print the outputs (embedding vectors) + for prompt, output in zip(prompts, outputs): + print("===============================") + print(f"Prompt: {prompt}\nEmbedding vector: {output['embedding']}") + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/examples/runtime/engine/fastapi_engine_inference.py b/examples/runtime/engine/fastapi_engine_inference.py new file mode 100644 index 000000000..a755cf8d8 --- /dev/null +++ b/examples/runtime/engine/fastapi_engine_inference.py @@ -0,0 +1,189 @@ +""" +FastAPI server example for text generation using SGLang Engine and demonstrating client usage. + +Starts the server, sends requests to it, and prints responses. + +Usage: +python fastapi_engine_inference.py --model-path Qwen/Qwen2.5-0.5B-Instruct --tp_size 1 --host 127.0.0.1 --port 8000 +""" + +import os +import subprocess +import time +from contextlib import asynccontextmanager + +import requests +from fastapi import FastAPI, Request + +import sglang as sgl +from sglang.utils import terminate_process + +engine = None + + +# Use FastAPI's lifespan manager to initialize/shutdown the engine +@asynccontextmanager +async def lifespan(app: FastAPI): + """Manages SGLang engine initialization during server startup.""" + global engine + # Initialize the SGLang engine when the server starts + # Adjust model_path and other engine arguments as needed + print("Loading SGLang engine...") + engine = sgl.Engine( + model_path=os.getenv("MODEL_PATH"), tp_size=int(os.getenv("TP_SIZE")) + ) + print("SGLang engine loaded.") + yield + # Clean up engine resources when the server stops (optional, depends on engine needs) + print("Shutting down SGLang engine...") + # engine.shutdown() # Or other cleanup if available/necessary + print("SGLang engine shutdown.") + + +app = FastAPI(lifespan=lifespan) + + +@app.post("/generate") +async def generate_text(request: Request): + """FastAPI endpoint to handle text generation requests.""" + global engine + if not engine: + return {"error": "Engine not initialized"}, 503 + + try: + data = await request.json() + prompt = data.get("prompt") + max_new_tokens = data.get("max_new_tokens", 128) + temperature = data.get("temperature", 0.7) + + if not prompt: + return {"error": "Prompt is required"}, 400 + + # Use async_generate for non-blocking generation + state = await engine.async_generate( + prompt, + sampling_params={ + "max_new_tokens": max_new_tokens, + "temperature": temperature, + }, + # Add other parameters like stop, top_p etc. as needed + ) + + return {"generated_text": state["text"]} + except Exception as e: + return {"error": str(e)}, 500 + + +# Helper function to start the server +def start_server(args, timeout=60): + """Starts the Uvicorn server as a subprocess and waits for it to be ready.""" + base_url = f"http://{args.host}:{args.port}" + command = [ + "python", + "-m", + "uvicorn", + "fastapi_engine_inference:app", + f"--host={args.host}", + f"--port={args.port}", + ] + + process = subprocess.Popen(command, stdout=None, stderr=None) + + start_time = time.perf_counter() + with requests.Session() as session: + while time.perf_counter() - start_time < timeout: + try: + # Check the /docs endpoint which FastAPI provides by default + response = session.get( + f"{base_url}/docs", timeout=5 + ) # Add a request timeout + if response.status_code == 200: + print(f"Server {base_url} is ready (responded on /docs)") + return process + except requests.ConnectionError: + # Specific exception for connection refused/DNS error etc. + pass + except requests.Timeout: + # Specific exception for request timeout + print(f"Health check to {base_url}/docs timed out, retrying...") + pass + except requests.RequestException as e: + # Catch other request exceptions + print(f"Health check request error: {e}, retrying...") + pass + # Use a shorter sleep interval for faster startup detection + time.sleep(1) + + # If loop finishes, raise the timeout error + # Attempt to terminate the failed process before raising + if process: + print( + "Server failed to start within timeout, attempting to terminate process..." + ) + terminate_process(process) # Use the imported terminate_process + raise TimeoutError( + f"Server failed to start at {base_url} within the timeout period." + ) + + +def send_requests(server_url, prompts, max_new_tokens, temperature): + """Sends generation requests to the running server for a list of prompts.""" + # Iterate through prompts and send requests + for i, prompt in enumerate(prompts): + print(f"\n[{i+1}/{len(prompts)}] Sending prompt: '{prompt}'") + payload = { + "prompt": prompt, + "max_new_tokens": max_new_tokens, + "temperature": temperature, + } + + try: + response = requests.post(f"{server_url}/generate", json=payload, timeout=60) + + result = response.json() + + print(f"Prompt: {prompt}\nResponse: {result['generated_text']}") + + except requests.exceptions.Timeout: + print(f" Error: Request timed out for prompt '{prompt}'") + except requests.exceptions.RequestException as e: + print(f" Error sending request for prompt '{prompt}': {e}") + + +if __name__ == "__main__": + """Main entry point for the script.""" + + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--model-path", type=str, default="Qwen/Qwen2.5-0.5B-Instruct") + parser.add_argument("--tp_size", type=int, default=1) + args = parser.parse_args() + + # Pass the model to the child uvicorn process via an env var + os.environ["MODEL_PATH"] = args.model_path + os.environ["TP_SIZE"] = str(args.tp_size) + + # Start the server + process = start_server(args) + + # Define the prompts and sampling parameters + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + max_new_tokens = 64 + temperature = 0.1 + + # Define server url + server_url = f"http://{args.host}:{args.port}" + + # Send requests to the server + send_requests(server_url, prompts, max_new_tokens, temperature) + + # Terminate the server process + terminate_process(process) diff --git a/examples/runtime/engine/launch_engine.py b/examples/runtime/engine/launch_engine.py new file mode 100644 index 000000000..b97e0bcd5 --- /dev/null +++ b/examples/runtime/engine/launch_engine.py @@ -0,0 +1,17 @@ +""" +This example demonstrates how to launch the offline engine. +""" + +import sglang as sgl + + +def main(): + llm = sgl.Engine(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") + llm.generate("What is the capital of France?") + llm.shutdown() + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/examples/runtime/engine/offline_batch_inference.py b/examples/runtime/engine/offline_batch_inference.py new file mode 100644 index 000000000..92e68dcd7 --- /dev/null +++ b/examples/runtime/engine/offline_batch_inference.py @@ -0,0 +1,43 @@ +""" +Usage: +python3 offline_batch_inference.py --model meta-llama/Llama-3.1-8B-Instruct +""" + +import argparse +import dataclasses + +import sglang as sgl +from sglang.srt.server_args import ServerArgs + + +def main( + server_args: ServerArgs, +): + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + # Create a sampling params object. + sampling_params = {"temperature": 0.8, "top_p": 0.95} + + # Create an LLM. + llm = sgl.Engine(**dataclasses.asdict(server_args)) + + outputs = llm.generate(prompts, sampling_params) + # Print the outputs. + for prompt, output in zip(prompts, outputs): + print("===============================") + print(f"Prompt: {prompt}\nGenerated text: {output['text']}") + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + main(server_args) diff --git a/examples/runtime/engine/offline_batch_inference_async.py b/examples/runtime/engine/offline_batch_inference_async.py new file mode 100644 index 000000000..578962d78 --- /dev/null +++ b/examples/runtime/engine/offline_batch_inference_async.py @@ -0,0 +1,65 @@ +""" +Usage: +python offline_batch_inference_async.py --model-path Qwen/Qwen2-VL-7B-Instruct + +Note: +This demo shows the usage of async generation, +which is useful to implement an online-like generation with batched inference. +""" + +import argparse +import asyncio +import dataclasses +import time + +import sglang as sgl +from sglang.srt.server_args import ServerArgs + + +class InferenceEngine: + def __init__(self, **kwargs): + self.engine = sgl.Engine(**kwargs) + + async def generate(self, prompt, sampling_params): + result = await self.engine.async_generate(prompt, sampling_params) + return result + + +async def run_server(server_args): + inference = InferenceEngine(**dataclasses.asdict(server_args)) + + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] * 100 + + # Create a sampling params object. + sampling_params = {"temperature": 0.8, "top_p": 0.95} + + # Run the generation tasks concurrently in async mode. + tasks = [] + for prompt in prompts: + task = asyncio.create_task(inference.generate(prompt, sampling_params)) + tasks.append(task) + + # Get and print the result + for task in tasks: + await task + while True: + if not task.done(): + time.sleep(1) + else: + result = task.result() + print(f"Generated text: {result['text']}") + break + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + asyncio.run(run_server(server_args)) diff --git a/examples/runtime/engine/offline_batch_inference_eagle.py b/examples/runtime/engine/offline_batch_inference_eagle.py new file mode 100644 index 000000000..a7a89ef5c --- /dev/null +++ b/examples/runtime/engine/offline_batch_inference_eagle.py @@ -0,0 +1,38 @@ +import sglang as sgl + + +def main(): + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + + # Create a sampling params object. + sampling_params = {"temperature": 0, "max_new_tokens": 30} + + # Create an LLM. + llm = sgl.Engine( + model_path="meta-llama/Llama-2-7b-chat-hf", + speculative_algorithm="EAGLE", + speculative_draft_model_path="lmsys/sglang-EAGLE-llama2-chat-7B", + speculative_num_steps=3, + speculative_eagle_topk=4, + speculative_num_draft_tokens=16, + cuda_graph_max_bs=8, + ) + + outputs = llm.generate(prompts, sampling_params) + + # Print the outputs. + for prompt, output in zip(prompts, outputs): + print("===============================") + print(f"Prompt: {prompt}\nGenerated text: {output['text']}") + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/examples/runtime/engine/offline_batch_inference_qwen_1m.py b/examples/runtime/engine/offline_batch_inference_qwen_1m.py new file mode 100644 index 000000000..664efa6d7 --- /dev/null +++ b/examples/runtime/engine/offline_batch_inference_qwen_1m.py @@ -0,0 +1,74 @@ +""" +Usage: +python3 offline_batch_inference.py +""" + +from urllib.request import urlopen + +import sglang as sgl + + +def load_prompt() -> str: + # Test cases with various lengths can be found at: + # + # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/64k.txt + # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/200k.txt + # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/600k.txt + # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt + + with urlopen( + "https://qianwen-res.oss-cn-beijing.aliyuncs.com" + "/Qwen2.5-1M/test-data/64k.txt", + timeout=5, + ) as response: + prompt = response.read().decode("utf-8") + return prompt + + +# Processing the prompt. +def process_requests(llm: sgl.Engine, prompts: list[str]) -> None: + # Create a sampling params object. + sampling_params = { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "repetition_penalty": 1.05, + "max_new_tokens": 256, + } + # Generate texts from the prompts. + outputs = llm.generate(prompts, sampling_params) + # Print the outputs. + for output in outputs: + prompt_token_ids = output["meta_info"]["prompt_tokens"] + generated_text = output["text"] + print( + f"Prompt length: {prompt_token_ids}, " f"Generated text: {generated_text!r}" + ) + + +# Create an LLM. +def initialize_engine() -> sgl.Engine: + llm = sgl.Engine( + model_path="Qwen/Qwen2.5-7B-Instruct-1M", + context_length=1048576, + page_size=256, + attention_backend="dual_chunk_flash_attn", + tp_size=4, + disable_radix_cache=True, + enable_mixed_chunk=False, + enable_torch_compile=False, + chunked_prefill_size=131072, + mem_fraction_static=0.6, + log_level="DEBUG", + ) + return llm + + +def main(): + llm = initialize_engine() + prompt = load_prompt() + process_requests(llm, [prompt]) + + +if __name__ == "__main__": + main() diff --git a/examples/runtime/engine/offline_batch_inference_vlm.py b/examples/runtime/engine/offline_batch_inference_vlm.py new file mode 100644 index 000000000..392823946 --- /dev/null +++ b/examples/runtime/engine/offline_batch_inference_vlm.py @@ -0,0 +1,52 @@ +""" +Usage: +python offline_batch_inference_vlm.py --model-path Qwen/Qwen2-VL-7B-Instruct +""" + +import argparse +import dataclasses + +import sglang as sgl +from sglang.srt.parser.conversation import chat_templates +from sglang.srt.server_args import ServerArgs + + +def main( + server_args: ServerArgs, +): + vlm = sgl.Engine(**dataclasses.asdict(server_args)) + + conv = chat_templates[server_args.chat_template].copy() + image_token = conv.image_token + + image_url = "https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true" + + prompt = f"What's in this image?\n{image_token}" + + sampling_params = { + "temperature": 0.001, + "max_new_tokens": 30, + } + + output = vlm.generate( + prompt=prompt, + image_data=image_url, + sampling_params=sampling_params, + ) + + print("===============================") + print(f"Prompt: {prompt}") + print(f"Generated text: {output['text']}") + + vlm.shutdown() + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + args = parser.parse_args() + + server_args = ServerArgs.from_cli_args(args) + main(server_args) diff --git a/examples/runtime/engine/readme.md b/examples/runtime/engine/readme.md new file mode 100644 index 000000000..fc6e89a6a --- /dev/null +++ b/examples/runtime/engine/readme.md @@ -0,0 +1,54 @@ +# SGLang Engine + +SGLang provides a direct inference engine without the need for an HTTP server. There are generally these use cases: + +- [Offline Batch Inference](#offline-batch-inference) +- [Embedding Generation](#embedding-generation) +- [Custom Server](#custom-server) +- [Token-In-Token-Out for RLHF](#token-in-token-out-for-rlhf) +- [Inference Using FastAPI](#inference-using-fastapi) + +## Examples + +### [Offline Batch Inference](./offline_batch_inference.py) + +In this example, we launch an SGLang engine and feed a batch of inputs for inference. If you provide a very large batch, the engine will intelligently schedule the requests to process efficiently and prevent OOM (Out of Memory) errors. + +### [Embedding Generation](./embedding.py) + +In this example, we launch an SGLang engine and feed a batch of inputs for embedding generation. + +### [Custom Server](./custom_server.py) + +This example demonstrates how to create a custom server on top of the SGLang Engine. We use [Sanic](https://sanic.dev/en/) as an example. The server supports both non-streaming and streaming endpoints. + +#### Steps + +1. Install Sanic: + + ```bash + pip install sanic + ``` + +2. Run the server: + + ```bash + python custom_server + ``` + +3. Send requests: + + ```bash + curl -X POST http://localhost:8000/generate -H "Content-Type: application/json" -d '{"prompt": "The Transformer architecture is..."}' + curl -X POST http://localhost:8000/generate_stream -H "Content-Type: application/json" -d '{"prompt": "The Transformer architecture is..."}' --no-buffer + ``` + + This will send both non-streaming and streaming requests to the server. + +### [Token-In-Token-Out for RLHF](../token_in_token_out) + +In this example, we launch an SGLang engine, feed tokens as input and generate tokens as output. + +### [Inference Using FastAPI](fastapi_engine_inference.py) + +This example demonstrates how to create a FastAPI server that uses the SGLang engine for text generation. diff --git a/examples/runtime/engine/save_remote_state.py b/examples/runtime/engine/save_remote_state.py new file mode 100644 index 000000000..a428195ca --- /dev/null +++ b/examples/runtime/engine/save_remote_state.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Saves each worker's model state dict directly to a checkpoint, which enables a +fast load path for large tensor-parallel models where each worker only needs to +read its own shard rather than the entire checkpoint. + +Example usage: + +python save_remote_state.py \ + --model-path /path/to/load \ + --tensor-parallel-size 8 \ + --remote-model-save-url [protocol]://[host]:[port]/[model_name] \ + +Then, the model can be loaded with + +llm = Engine( + model_path="[protocol]://[host]:[port]/[model_name]", + tensor_parallel_size=8, +) +""" +import dataclasses +from argparse import ArgumentParser +from pathlib import Path + +from sglang import Engine, ServerArgs + +parser = ArgumentParser() +ServerArgs.add_cli_args(parser) + +parser.add_argument( + "--remote-model-save-url", + required=True, + type=str, + help="remote address to store model weights", +) +parser.add_argument( + "--remote-draft-model-save-url", + default=None, + type=str, + help="remote address to store draft model weights", +) + + +def main(args): + engine_args = ServerArgs.from_cli_args(args) + model_path = engine_args.model_path + if not Path(model_path).is_dir(): + raise ValueError("model path must be a local directory") + # Create LLM instance from arguments + llm = Engine(**dataclasses.asdict(engine_args)) + llm.save_remote_model( + url=args.remote_model_save_url, draft_url=args.remote_draft_model_save_url + ) + print("save remote (draft) model successfully") + + +if __name__ == "__main__": + args = parser.parse_args() + main(args) diff --git a/examples/runtime/engine/save_sharded_state.py b/examples/runtime/engine/save_sharded_state.py new file mode 100644 index 000000000..80ad5321f --- /dev/null +++ b/examples/runtime/engine/save_sharded_state.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Saves each worker's model state dict directly to a checkpoint, which enables a +fast load path for large tensor-parallel models where each worker only needs to +read its own shard rather than the entire checkpoint. + +Example usage: + +python save_sharded_state.py \ + --model-path /path/to/load \ + --quantization deepspeedfp \ + --tensor-parallel-size 8 \ + --output /path/to/save + +Then, the model can be loaded with + +llm = Engine( + model_path="/path/to/save", + load_format="sharded_state", + quantization="deepspeedfp", + tensor_parallel_size=8, +) +""" +import dataclasses +import os +import shutil +from argparse import ArgumentParser +from pathlib import Path + +from sglang import Engine, ServerArgs + +parser = ArgumentParser() +ServerArgs.add_cli_args(parser) + +parser.add_argument( + "--output", "-o", required=True, type=str, help="path to output checkpoint" +) +parser.add_argument( + "--file-pattern", type=str, help="string pattern of saved filenames" +) +parser.add_argument( + "--max-file-size", + type=str, + default=5 * 1024**3, + help="max size (in bytes) of each safetensors file", +) + + +def main(args): + engine_args = ServerArgs.from_cli_args(args) + model_path = engine_args.model_path + if not Path(model_path).is_dir(): + raise ValueError("model path must be a local directory") + # Create LLM instance from arguments + llm = Engine(**dataclasses.asdict(engine_args)) + Path(args.output).mkdir(exist_ok=True) + llm.save_sharded_model( + path=args.output, pattern=args.file_pattern, max_size=args.max_file_size + ) + + # Copy metadata files to output directory + for file in os.listdir(model_path): + if os.path.splitext(file)[1] not in (".bin", ".pt", ".safetensors"): + if os.path.isdir(os.path.join(model_path, file)): + shutil.copytree( + os.path.join(model_path, file), os.path.join(args.output, file) + ) + else: + shutil.copy(os.path.join(model_path, file), args.output) + + +if __name__ == "__main__": + args = parser.parse_args() + main(args) diff --git a/examples/runtime/hidden_states/hidden_states_engine.py b/examples/runtime/hidden_states/hidden_states_engine.py new file mode 100644 index 000000000..60ab302ca --- /dev/null +++ b/examples/runtime/hidden_states/hidden_states_engine.py @@ -0,0 +1,66 @@ +""" +Usage: +python hidden_states.py + +Note that each time you change the `return_hidden_states` parameter, +the cuda graph will be recaptured, which might lead to a performance hit. +So avoid getting hidden states and completions alternately. +""" + +import torch + +import sglang as sgl + + +def main(): + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + # Create an LLM. + llm = sgl.Engine( + model_path="Alibaba-NLP/gte-Qwen2-1.5B-instruct", + enable_return_hidden_states=True, + ) + + sampling_params = { + "temperature": 0.8, + "top_p": 0.95, + "max_new_tokens": 10, + } + + outputs = llm.generate( + prompts, sampling_params=sampling_params, return_hidden_states=True + ) + + llm.shutdown() + + for prompt, output in zip(prompts, outputs): + for i in range(len(output["meta_info"]["hidden_states"])): + output["meta_info"]["hidden_states"][i] = torch.tensor( + output["meta_info"]["hidden_states"][i], dtype=torch.bfloat16 + ) + print("===============================") + print( + f"Prompt: {prompt}\n" + f"Generated text: {output['text']}\n" + f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t" + f"Completion_tokens: {output['meta_info']['completion_tokens']}" + ) + print("Hidden states: ") + hidden_states = torch.cat( + [ + i.unsqueeze(0) if len(i.shape) == 1 else i + for i in output["meta_info"]["hidden_states"] + ] + ) + print(hidden_states) + print() + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/examples/runtime/hidden_states/hidden_states_server.py b/examples/runtime/hidden_states/hidden_states_server.py new file mode 100644 index 000000000..b04f74372 --- /dev/null +++ b/examples/runtime/hidden_states/hidden_states_server.py @@ -0,0 +1,81 @@ +""" +Usage: + +python hidden_states_server.py + +Note that each time you change the `return_hidden_states` parameter, +the cuda graph will be recaptured, which might lead to a performance hit. +So avoid getting hidden states and completions alternately. +""" + +import requests +import torch + +from sglang.test.test_utils import is_in_ci +from sglang.utils import terminate_process, wait_for_server + +if is_in_ci(): + from docs.backend.patch import launch_server_cmd +else: + from sglang.utils import launch_server_cmd + + +def main(): + # Launch the server + server_process, port = launch_server_cmd( + "python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct --enable-return-hidden-states --host 0.0.0.0" + ) + wait_for_server(f"http://localhost:{port}") + + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + + sampling_params = { + "temperature": 0.8, + "top_p": 0.95, + "max_new_tokens": 10, + } + + json_data = { + "text": prompts, + "sampling_params": sampling_params, + "return_hidden_states": True, + } + + response = requests.post( + f"http://localhost:{port}/generate", + json=json_data, + ) + + terminate_process(server_process) + + outputs = response.json() + for prompt, output in zip(prompts, outputs): + for i in range(len(output["meta_info"]["hidden_states"])): + output["meta_info"]["hidden_states"][i] = torch.tensor( + output["meta_info"]["hidden_states"][i], dtype=torch.bfloat16 + ) + print("===============================") + print( + f"Prompt: {prompt}\n" + f"Generated text: {output['text']}\n" + f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t" + f"Completion_tokens: {output['meta_info']['completion_tokens']}" + ) + print("Hidden states: ") + hidden_states = torch.cat( + [ + i.unsqueeze(0) if len(i.shape) == 1 else i + for i in output["meta_info"]["hidden_states"] + ] + ) + print(hidden_states) + print() + + +if __name__ == "__main__": + main() diff --git a/examples/runtime/lora.py b/examples/runtime/lora.py new file mode 100644 index 000000000..bf3fc2d9e --- /dev/null +++ b/examples/runtime/lora.py @@ -0,0 +1,37 @@ +# launch server +# python -m sglang.launch_server --model mistralai/Mistral-7B-Instruct-v0.3 --lora-paths /home/ying/test_lora lora1=/home/ying/test_lora_1 lora2=/home/ying/test_lora_2 --disable-radix --disable-cuda-graph --max-loras-per-batch 4 + +# send requests +# lora_path[i] specifies the LoRA used for text[i], so make sure they have the same length +# use None to specify base-only prompt, e.x. "lora_path": [None, "/home/ying/test_lora"] +import json + +import requests + +url = "http://127.0.0.1:30000" +json_data = { + "text": [ + "prompt 1", + "prompt 2", + "prompt 3", + "prompt 4", + "prompt 5", + "prompt 6", + "prompt 7", + ], + "sampling_params": {"max_new_tokens": 32}, + "lora_path": [ + "/home/ying/test_lora", + "lora1", + "lora2", + "lora1", + "lora2", + None, + None, + ], +} +response = requests.post( + url + "/generate", + json=json_data, +) +print(json.dumps(response.json())) diff --git a/examples/runtime/multimodal/llama3_llava_server.py b/examples/runtime/multimodal/llama3_llava_server.py new file mode 100644 index 000000000..a8409af71 --- /dev/null +++ b/examples/runtime/multimodal/llama3_llava_server.py @@ -0,0 +1,111 @@ +""" +Usage: +# Installing latest llava-next: pip install git+https://github.com/LLaVA-VL/LLaVA-NeXT.git +# Installing latest sglang. + +# Endpoint Service CLI: +python -m sglang.launch_server --model-path lmms-lab/llama3-llava-next-8b --port=30000 + +python3 llama3_llava_server.py + +Output: +"Friends posing for a fun photo with a life-sized teddy bear, creating a playful and memorable moment." +""" + +import argparse +import asyncio +import copy +import json + +import aiohttp +import requests +from llava.conversation import conv_llava_llama_3 + + +async def send_request(url, data, delay=0): + await asyncio.sleep(delay) + async with aiohttp.ClientSession() as session: + async with session.post(url, json=data) as resp: + output = await resp.json() + return output + + +async def test_concurrent(args): + url = f"{args.host}:{args.port}" + + prompt = "\nPlease generate caption towards this image." + conv_template = copy.deepcopy(conv_llava_llama_3) + conv_template.append_message(role=conv_template.roles[0], message=prompt) + conv_template.append_message(role=conv_template.roles[1], message=None) + prompt_with_template = conv_template.get_prompt() + response = [] + for i in range(1): + response.append( + send_request( + url + "/generate", + { + "text": prompt_with_template, + "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", + "sampling_params": { + "max_new_tokens": 1024, + "temperature": 0, + "top_p": 1.0, + "presence_penalty": 2, + "frequency_penalty": 2, + "stop": "<|eot_id|>", + }, + }, + ) + ) + + rets = await asyncio.gather(*response) + for ret in rets: + print(ret["text"]) + + +def test_streaming(args): + url = f"{args.host}:{args.port}" + prompt = "\nPlease generate caption towards this image." + conv_template = copy.deepcopy(conv_llava_llama_3) + conv_template.append_message(role=conv_template.roles[0], message=prompt) + conv_template.append_message(role=conv_template.roles[1], message=None) + prompt_with_template = conv_template.get_prompt() + pload = { + "text": prompt_with_template, + "sampling_params": { + "max_new_tokens": 1024, + "temperature": 0, + "top_p": 1.0, + "presence_penalty": 2, + "frequency_penalty": 2, + "stop": "<|eot_id|>", + }, + "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", + "stream": True, + } + response = requests.post( + url + "/generate", + json=pload, + stream=True, + ) + + prev = 0 + for chunk in response.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + output = data["text"].strip() + print(output[prev:], end="", flush=True) + prev = len(output) + print("") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="http://127.0.0.1") + parser.add_argument("--port", type=int, default=30000) + args = parser.parse_args() + asyncio.run(test_concurrent(args)) + test_streaming(args) diff --git a/examples/runtime/multimodal/llava_onevision_server.py b/examples/runtime/multimodal/llava_onevision_server.py new file mode 100644 index 000000000..ee921b558 --- /dev/null +++ b/examples/runtime/multimodal/llava_onevision_server.py @@ -0,0 +1,264 @@ +""" +Usage: + +python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-72b-ov --port=30000 --tp-size=8 + +python3 llava_onevision_server.py +""" + +import base64 +import io +import os +import sys +import time + +import numpy as np +import openai +import requests +from decord import VideoReader, cpu +from PIL import Image + +# pip install httpx==0.23.3 +# pip install decord +# pip install protobuf==3.20.0 + + +def download_video(url, cache_dir): + file_path = os.path.join(cache_dir, "jobs.mp4") + os.makedirs(cache_dir, exist_ok=True) + + response = requests.get(url) + response.raise_for_status() + + with open(file_path, "wb") as f: + f.write(response.content) + + print(f"File downloaded and saved to: {file_path}") + return file_path + + +def create_openai_client(base_url): + return openai.Client(api_key="EMPTY", base_url=base_url) + + +def image_stream_request_test(client): + print("----------------------Image Stream Request Test----------------------") + stream_request = client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" + }, + }, + { + "type": "text", + "text": "Please describe this image. Please list the benchmarks and the models.", + }, + ], + }, + ], + temperature=0.7, + max_tokens=1024, + stream=True, + ) + stream_response = "" + + for chunk in stream_request: + if chunk.choices[0].delta.content is not None: + content = chunk.choices[0].delta.content + stream_response += content + sys.stdout.write(content) + sys.stdout.flush() + + print("-" * 30) + + +def multi_image_stream_request_test(client): + print( + "----------------------Multi-Images Stream Request Test----------------------" + ) + stream_request = client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" + }, + "modalities": "multi-images", + }, + { + "type": "image_url", + "image_url": { + "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/test/lang/example_image.png" + }, + "modalities": "multi-images", + }, + { + "type": "text", + "text": "I have shown you two images. Please describe the two images to me.", + }, + ], + }, + ], + temperature=0.7, + max_tokens=1024, + stream=True, + ) + stream_response = "" + + for chunk in stream_request: + if chunk.choices[0].delta.content is not None: + content = chunk.choices[0].delta.content + stream_response += content + sys.stdout.write(content) + sys.stdout.flush() + + print("-" * 30) + + +def video_stream_request_test(client, video_path): + print("------------------------Video Stream Request Test----------------------") + messages = prepare_video_messages(video_path) + + video_request = client.chat.completions.create( + model="default", + messages=messages, + temperature=0, + max_tokens=1024, + stream=True, + ) + print("-" * 30) + video_response = "" + + for chunk in video_request: + if chunk.choices[0].delta.content is not None: + content = chunk.choices[0].delta.content + video_response += content + sys.stdout.write(content) + sys.stdout.flush() + print("-" * 30) + + +def image_speed_test(client): + print("----------------------Image Speed Test----------------------") + start_time = time.perf_counter() + request = client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" + }, + }, + { + "type": "text", + "text": "Please describe this image. Please list the benchmarks and the models.", + }, + ], + }, + ], + temperature=0, + max_tokens=1024, + ) + end_time = time.perf_counter() + response = request.choices[0].message.content + print(response) + print("-" * 30) + print_speed_test_results(request, start_time, end_time) + + +def video_speed_test(client, video_path): + print("------------------------Video Speed Test------------------------") + messages = prepare_video_messages(video_path) + + start_time = time.perf_counter() + video_request = client.chat.completions.create( + model="default", + messages=messages, + temperature=0, + max_tokens=1024, + ) + end_time = time.perf_counter() + video_response = video_request.choices[0].message.content + print(video_response) + print("-" * 30) + print_speed_test_results(video_request, start_time, end_time) + + +def prepare_video_messages(video_path): + max_frames_num = 32 + vr = VideoReader(video_path, ctx=cpu(0)) + total_frame_num = len(vr) + uniform_sampled_frames = np.linspace( + 0, total_frame_num - 1, max_frames_num, dtype=int + ) + frame_idx = uniform_sampled_frames.tolist() + frames = vr.get_batch(frame_idx).asnumpy() + + base64_frames = [] + for frame in frames: + pil_img = Image.fromarray(frame) + buff = io.BytesIO() + pil_img.save(buff, format="JPEG") + base64_str = base64.b64encode(buff.getvalue()).decode("utf-8") + base64_frames.append(base64_str) + + messages = [{"role": "user", "content": []}] + + for base64_frame in base64_frames: + frame_format = { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{base64_frame}"}, + "modalities": "video", + } + messages[0]["content"].append(frame_format) + + prompt = {"type": "text", "text": "Please describe the video in detail."} + messages[0]["content"].append(prompt) + + return messages + + +def print_speed_test_results(request, start_time, end_time): + total_tokens = request.usage.total_tokens + completion_tokens = request.usage.completion_tokens + prompt_tokens = request.usage.prompt_tokens + + print(f"Total tokens: {total_tokens}") + print(f"Completion tokens: {completion_tokens}") + print(f"Prompt tokens: {prompt_tokens}") + print(f"Time taken: {end_time - start_time} seconds") + print(f"Token per second: {total_tokens / (end_time - start_time)}") + print(f"Completion token per second: {completion_tokens / (end_time - start_time)}") + print(f"Prompt token per second: {prompt_tokens / (end_time - start_time)}") + + +def main(): + url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4" + cache_dir = os.path.expanduser("~/.cache") + video_path = download_video(url, cache_dir) + + client = create_openai_client("http://127.0.0.1:30000/v1") + + image_stream_request_test(client) + multi_image_stream_request_test(client) + video_stream_request_test(client, video_path) + image_speed_test(client) + video_speed_test(client, video_path) + + +if __name__ == "__main__": + main() diff --git a/examples/runtime/multimodal/pixtral_server.py b/examples/runtime/multimodal/pixtral_server.py new file mode 100644 index 000000000..d907de14d --- /dev/null +++ b/examples/runtime/multimodal/pixtral_server.py @@ -0,0 +1,127 @@ +""" +Usage: +# Run a Pixtral model with SGLang: +# HuggingFace: +python -m sglang.launch_server --model-path mistral-community/pixtral-12b --port=30000 +# ModelScope: +python -m sglang.launch_server --model-path AI-ModelScope/pixtral-12b --port=30000 + +# Then test it with: +python pixtral_server.py + +This script tests Pixtral model with both single and multiple images. +""" + +import argparse +import asyncio +import json + +import aiohttp +import requests + +IMAGE_TOKEN_SEP = "\n[IMG]" +ROUTE = "/generate" + + +async def send_request(url, data, delay=0): + await asyncio.sleep(delay) + async with aiohttp.ClientSession() as session: + async with session.post(url, json=data) as resp: + output = await resp.json() + return output + + +async def test_concurrent(args): + url = f"{args.host}:{args.port}{ROUTE}" + + # Single image test + if args.single_image: + prompt = f"[INST]Describe this image in detail.{IMAGE_TOKEN_SEP}[/INST]" + image_url = "https://picsum.photos/id/237/400/300" + modality = ["image"] + # Multiple images test + else: + image_urls = [ + "https://picsum.photos/id/237/400/300", + "https://picsum.photos/id/27/500/500", + ] + prompt = f"[INST]How many photos are there? Describe each in a very short sentence.{IMAGE_TOKEN_SEP * len(image_urls)}[/INST]" + image_url = image_urls + modality = ["multi-images"] + + response = await send_request( + url, + { + "text": prompt, + "image_data": image_url, + "sampling_params": { + "max_new_tokens": 100, + "temperature": 0.7, + "top_p": 0.9, + }, + "modalities": modality, + }, + ) + + print(f"Response: {response}") + if "text" in response: + print("\nOutput text:", response["text"]) + + +def test_streaming(args): + url = f"{args.host}:{args.port}/generate" + + # Single image test + if args.single_image: + prompt = f"[INST]Describe this image in detail.{IMAGE_TOKEN_SEP}[/INST]" + image_data = "https://picsum.photos/id/237/400/300" + modality = ["image"] + # Multiple images test + else: + image_urls = [ + "https://picsum.photos/id/237/400/300", + "https://picsum.photos/id/27/500/500", + ] + prompt = f"[INST]How many photos are there? Describe each in a very short sentence.{IMAGE_TOKEN_SEP * len(image_urls)}[/INST]" + image_data = image_urls + modality = ["multi-images"] + + pload = { + "text": prompt, + "image_data": image_data, + "sampling_params": {"max_new_tokens": 100, "temperature": 0.7, "top_p": 0.9}, + "modalities": modality, + "stream": True, + } + + response = requests.post(url, json=pload, stream=True) + + print("Streaming response:") + prev = 0 + for chunk in response.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + output = data["text"].strip() + print(output[prev:], end="", flush=True) + prev = len(output) + print("\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="http://127.0.0.1") + parser.add_argument("--port", type=int, default=30000) + parser.add_argument( + "--single-image", + action="store_true", + help="Test with single image instead of multiple images", + ) + parser.add_argument("--no-stream", action="store_true", help="Don't test streaming") + args = parser.parse_args() + + asyncio.run(test_concurrent(args)) + if not args.no_stream: + test_streaming(args) diff --git a/examples/runtime/multimodal/qwen_llava_server.py b/examples/runtime/multimodal/qwen_llava_server.py new file mode 100644 index 000000000..d8b3226e7 --- /dev/null +++ b/examples/runtime/multimodal/qwen_llava_server.py @@ -0,0 +1,111 @@ +""" +Usage: +# Installing latest llava-next: pip install git+https://github.com/LLaVA-VL/LLaVA-NeXT.git +# Installing latest sglang. + +# Endpoint Service CLI: +python -m sglang.launch_server --model-path lmms-lab/llava-next-72b --port=30000 --tp-size=8 + +python3 qwen_llava_server.py + +Output: +"Two children pose with a large teddy bear, one holding a smaller stuffed bear, in a room with an American flag and potted plants." +""" + +import argparse +import asyncio +import copy +import json + +import aiohttp +import requests +from llava.conversation import conv_qwen + + +async def send_request(url, data, delay=0): + await asyncio.sleep(delay) + async with aiohttp.ClientSession() as session: + async with session.post(url, json=data) as resp: + output = await resp.json() + return output + + +async def test_concurrent(args): + url = f"{args.host}:{args.port}" + + prompt = "\nPlease generate caption towards this image." + conv_template = copy.deepcopy(conv_qwen) + conv_template.append_message(role=conv_template.roles[0], message=prompt) + conv_template.append_message(role=conv_template.roles[1], message=None) + prompt_with_template = conv_template.get_prompt() + response = [] + for i in range(1): + response.append( + send_request( + url + "/generate", + { + "text": prompt_with_template, + "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", + "sampling_params": { + "max_new_tokens": 1024, + "temperature": 0, + "top_p": 1.0, + "presence_penalty": 2, + "frequency_penalty": 2, + "stop": "<|im_end|>", + }, + }, + ) + ) + + rets = await asyncio.gather(*response) + for ret in rets: + print(ret["text"]) + + +def test_streaming(args): + url = f"{args.host}:{args.port}" + prompt = "\nPlease generate caption towards this image." + conv_template = copy.deepcopy(conv_qwen) + conv_template.append_message(role=conv_template.roles[0], message=prompt) + conv_template.append_message(role=conv_template.roles[1], message=None) + prompt_with_template = conv_template.get_prompt() + pload = { + "text": prompt_with_template, + "sampling_params": { + "max_new_tokens": 1024, + "temperature": 0, + "top_p": 1.0, + "presence_penalty": 2, + "frequency_penalty": 2, + "stop": "<|im_end|>", + }, + "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", + "stream": True, + } + response = requests.post( + url + "/generate", + json=pload, + stream=True, + ) + + prev = 0 + for chunk in response.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + output = data["text"].strip() + print(output[prev:], end="", flush=True) + prev = len(output) + print("") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="http://127.0.0.1") + parser.add_argument("--port", type=int, default=30000) + args = parser.parse_args() + asyncio.run(test_concurrent(args)) + test_streaming(args) diff --git a/examples/runtime/multimodal_embedding.py b/examples/runtime/multimodal_embedding.py new file mode 100644 index 000000000..4e8d748b4 --- /dev/null +++ b/examples/runtime/multimodal_embedding.py @@ -0,0 +1,18 @@ +# launch server +# python -m sglang.launch_server --model-path Alibaba-NLP/gme-Qwen2-VL-2B-Instruct --is-embedding + +import requests + +url = "http://127.0.0.1:30000" + +text_input = "Represent this image in embedding space." +image_path = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/023.jpg" + +payload = { + "model": "gme-qwen2-vl", + "input": [{"text": text_input}, {"image": image_path}], +} + +response = requests.post(url + "/v1/embeddings", json=payload).json() + +print("Embeddings:", [x.get("embedding") for x in response.get("data", [])]) diff --git a/examples/runtime/openai_chat_with_response_prefill.py b/examples/runtime/openai_chat_with_response_prefill.py new file mode 100644 index 000000000..6d803a1d1 --- /dev/null +++ b/examples/runtime/openai_chat_with_response_prefill.py @@ -0,0 +1,53 @@ +""" +Usage: +1) Launch the server in one terminal: + python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --port 30000 + +2) Run this script in another terminal: + python openai_chat_with_response_prefill.py + +This example demonstrates two chat completion calls: +- One with continue_final_message enabled (the final assistant message is used as a prefill). +- One without continue_final_message (the final assistant message remains, starting a new turn). +""" + +import openai + +client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") + +messages = [ + {"role": "system", "content": "You are a helpful AI assistant."}, + { + "role": "user", + "content": """ +Extract the name, size, price, and color from this product description as a JSON object: + + +The SmartHome Mini is a compact smart home assistant available in black or white for only $49.99. +At just 5 inches wide, it lets you control lights, thermostats, and other connected devices via voice or app— +no matter where you place it in your home. +This affordable little hub brings convenient hands-free control to your smart devices. + +""", + }, + {"role": "assistant", "content": "{\n"}, +] + +# Calling the API with continue_final_message enabled. +print("=== Prefill with continue_final_messagem ===") +response_with = client.chat.completions.create( + model="meta-llama/Llama-3.1-8B-Instruct", + messages=messages, + temperature=0, + extra_body={"continue_final_message": True}, +) +print(response_with.choices[0].message.content) + +# Calling the API without continue_final_message (using default behavior). +print("\n=== Prefill without continue_final_message ===") +response_without = client.chat.completions.create( + model="meta-llama/Llama-3.1-8B-Instruct", + messages=messages, + temperature=0, +) +print(response_without.choices[0].message.content) diff --git a/examples/runtime/reward_model.py b/examples/runtime/reward_model.py new file mode 100644 index 000000000..1a1177e66 --- /dev/null +++ b/examples/runtime/reward_model.py @@ -0,0 +1,32 @@ +# launch server +# python -m sglang.launch_server --model LxzGordon/URM-LLaMa-3.1-8B --is-embedding + +import requests + +url = "http://127.0.0.1:30000" + +PROMPT = ( + "What is the range of the numeric output of a sigmoid node in a neural network?" +) +RESPONSE1 = "The output of a sigmoid node is bounded between -1 and 1." +RESPONSE2 = "The output of a sigmoid node is bounded between 0 and 1." + +json_data = { + "conv": [ + [ + {"role": "user", "content": PROMPT}, + {"role": "assistant", "content": RESPONSE1}, + ], + [ + {"role": "user", "content": PROMPT}, + {"role": "assistant", "content": RESPONSE2}, + ], + ], +} +response = requests.post( + url + "/classify", + json=json_data, +).json() + +print(response) +print("scores:", [x["embedding"] for x in response]) diff --git a/examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py b/examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py new file mode 100644 index 000000000..cb1b7ddc1 --- /dev/null +++ b/examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py @@ -0,0 +1,43 @@ +""" +This example demonstrates how to provide tokenized ids to LLM as input instead of text prompt, i.e. a token-in-token-out workflow. +""" + +import sglang as sgl +from sglang.srt.hf_transformers_utils import get_tokenizer + +MODEL_PATH = "meta-llama/Llama-3.1-8B-Instruct" + + +def main(): + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + # Create a sampling params object. + sampling_params = {"temperature": 0.8, "top_p": 0.95} + + # Tokenize inputs + tokenizer = get_tokenizer(MODEL_PATH) + token_ids_list = [tokenizer.encode(prompt) for prompt in prompts] + + # Create an LLM. + llm = sgl.Engine(model_path=MODEL_PATH, skip_tokenizer_init=True) + + outputs = llm.generate(input_ids=token_ids_list, sampling_params=sampling_params) + # Print the outputs. + for prompt, output in zip(prompts, outputs): + decode_output = tokenizer.decode(output["output_ids"]) + print("===============================") + print( + f"Prompt: {prompt}\nGenerated token ids: {output['output_ids']}\nGenerated text: {decode_output}" + ) + print() + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/examples/runtime/token_in_token_out/token_in_token_out_llm_server.py b/examples/runtime/token_in_token_out/token_in_token_out_llm_server.py new file mode 100644 index 000000000..00c0988b2 --- /dev/null +++ b/examples/runtime/token_in_token_out/token_in_token_out_llm_server.py @@ -0,0 +1,68 @@ +""" +Usage: + +python token_in_token_out_llm_server.py + +""" + +import requests + +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.test.test_utils import is_in_ci +from sglang.utils import terminate_process, wait_for_server + +if is_in_ci(): + from docs.backend.patch import launch_server_cmd +else: + from sglang.utils import launch_server_cmd + + +MODEL_PATH = "meta-llama/Llama-3.1-8B-Instruct" + + +def main(): + # Launch the server + server_process, port = launch_server_cmd( + f"python -m sglang.launch_server --model-path {MODEL_PATH} --skip-tokenizer-init --host 0.0.0.0" + ) + wait_for_server(f"http://localhost:{port}") + + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + + # Create a sampling params object. + sampling_params = {"temperature": 0.8, "top_p": 0.95} + + # Tokenize inputs + tokenizer = get_tokenizer(MODEL_PATH) + token_ids_list = [tokenizer.encode(prompt) for prompt in prompts] + + json_data = { + "input_ids": token_ids_list, + "sampling_params": sampling_params, + } + + response = requests.post( + f"http://localhost:{port}/generate", + json=json_data, + ) + + outputs = response.json() + for prompt, output in zip(prompts, outputs): + print("===============================") + decode_output = tokenizer.decode(output["output_ids"]) + print( + f"Prompt: {prompt}\nGenerated token ids: {output['output_ids']}\nGenerated text: {decode_output}" + ) + print() + + terminate_process(server_process) + + +if __name__ == "__main__": + main() diff --git a/examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py b/examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py new file mode 100644 index 000000000..bf6f19719 --- /dev/null +++ b/examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py @@ -0,0 +1,74 @@ +import argparse +import dataclasses +from typing import Tuple + +from transformers import AutoProcessor + +from sglang import Engine +from sglang.lang.chat_template import get_chat_template_by_model_path +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.server_args import ServerArgs +from sglang.test.test_utils import DEFAULT_IMAGE_URL + + +def get_input_ids( + server_args: ServerArgs, model_config: ModelConfig +) -> Tuple[list[int], list]: + chat_template = get_chat_template_by_model_path(model_config.model_path) + text = f"{chat_template.image_token}What is in this picture?" + image_data = [DEFAULT_IMAGE_URL] + + processor = AutoProcessor.from_pretrained( + model_config.model_path, trust_remote_code=server_args.trust_remote_code + ) + + input_ids = ( + processor.tokenizer( + text=[text], + return_tensors="pt", + ) + .input_ids[0] + .tolist() + ) + + return input_ids, image_data + + +def token_in_out_example( + server_args: ServerArgs, +): + input_ids, image_data = get_input_ids( + server_args, + ModelConfig( + server_args.model_path, + trust_remote_code=server_args.trust_remote_code, + model_override_args=server_args.json_model_override_args, + ), + ) + backend = Engine(**dataclasses.asdict(server_args)) + + output = backend.generate( + input_ids=input_ids, + image_data=image_data, + sampling_params={ + "temperature": 0.8, + "max_new_tokens": 32, + }, + ) + + print("===============================") + print(f"Output token ids: ", output["output_ids"]) + + backend.shutdown() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + args = [ + "--model-path=Qwen/Qwen2-VL-2B", + ] + args = parser.parse_args(args=args) + server_args = ServerArgs.from_cli_args(args) + server_args.skip_tokenizer_init = True + token_in_out_example(server_args) diff --git a/examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py b/examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py new file mode 100644 index 000000000..392e1bf0e --- /dev/null +++ b/examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py @@ -0,0 +1,78 @@ +""" +Usage: + +python token_in_token_out_vlm_server.py + +""" + +from typing import Tuple + +import requests +from transformers import AutoProcessor + +from sglang.lang.chat_template import get_chat_template_by_model_path +from sglang.test.test_utils import DEFAULT_IMAGE_URL, is_in_ci +from sglang.utils import terminate_process, wait_for_server + +if is_in_ci(): + from docs.backend.patch import launch_server_cmd +else: + from sglang.utils import launch_server_cmd + + +MODEL_PATH = "Qwen/Qwen2-VL-2B" + + +def get_input_ids() -> Tuple[list[int], list]: + chat_template = get_chat_template_by_model_path(MODEL_PATH) + text = f"{chat_template.image_token}What is in this picture?" + image_data = [DEFAULT_IMAGE_URL] + + processor = AutoProcessor.from_pretrained(MODEL_PATH) + + input_ids = ( + processor.tokenizer( + text=[text], + return_tensors="pt", + ) + .input_ids[0] + .tolist() + ) + + return input_ids, image_data + + +def main(): + # Launch the server + server_process, port = launch_server_cmd( + f"python -m sglang.launch_server --model-path {MODEL_PATH} --skip-tokenizer-init --host 0.0.0.0" + ) + wait_for_server(f"http://localhost:{port}") + + input_ids, image_data = get_input_ids() + + sampling_params = { + "temperature": 0.8, + "max_new_tokens": 32, + } + + json_data = { + "input_ids": input_ids, + "image_data": image_data, + "sampling_params": sampling_params, + } + + response = requests.post( + f"http://localhost:{port}/generate", + json=json_data, + ) + + output = response.json() + print("===============================") + print(f"Output token ids: ", output["output_ids"]) + + terminate_process(server_process) + + +if __name__ == "__main__": + main() diff --git a/examples/runtime/vertex_predict.py b/examples/runtime/vertex_predict.py new file mode 100644 index 000000000..58a41b1c4 --- /dev/null +++ b/examples/runtime/vertex_predict.py @@ -0,0 +1,66 @@ +""" +Usage: +python -m sglang.launch_server --model meta-llama/Llama-2-7b-hf --port 30000 +python vertex_predict.py + +This example shows the request and response formats of the prediction route for +Google Cloud Vertex AI Online Predictions. + +Vertex AI SDK for Python is recommended for deploying models to Vertex AI +instead of a local server. After deploying the model to a Vertex AI Online +Prediction Endpoint, send requests via the Python SDK: + +response = endpoint.predict( + instances=[ + {"text": "The capital of France is"}, + {"text": "What is a car?"}, + ], + parameters={"sampling_params": {"max_new_tokens": 16}}, +) +print(response.predictions) + +More details about get online predictions from Vertex AI can be found at +https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions. +""" + +from dataclasses import dataclass +from typing import List, Optional + +import requests + + +@dataclass +class VertexPrediction: + predictions: List + + +class LocalVertexEndpoint: + def __init__(self) -> None: + self.base_url = "http://127.0.0.1:30000" + + def predict(self, instances: List[dict], parameters: Optional[dict] = None): + response = requests.post( + self.base_url + "/vertex_generate", + json={ + "instances": instances, + "parameters": parameters, + }, + ) + return VertexPrediction(predictions=response.json()["predictions"]) + + +endpoint = LocalVertexEndpoint() + +# Predict with a single prompt. +response = endpoint.predict(instances=[{"text": "The capital of France is"}]) +print(response.predictions) + +# Predict with multiple prompts and parameters. +response = endpoint.predict( + instances=[ + {"text": "The capital of France is"}, + {"text": "What is a car?"}, + ], + parameters={"sampling_params": {"max_new_tokens": 16}}, +) +print(response.predictions) diff --git a/package-lock.json b/package-lock.json new file mode 100644 index 000000000..54cb66f0b --- /dev/null +++ b/package-lock.json @@ -0,0 +1,6 @@ +{ + "name": "sglang", + "lockfileVersion": 3, + "requires": true, + "packages": {} +} diff --git a/python/pyproject.toml b/python/pyproject.toml new file mode 100755 index 000000000..f2e69b3c0 --- /dev/null +++ b/python/pyproject.toml @@ -0,0 +1,167 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "sglang" +version = "0.5.2" +description = "SGLang is a fast serving framework for large language models and vision language models." +readme = "README.md" +requires-python = ">=3.10" +license = { file = "LICENSE" } +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", +] +dependencies = ["aiohttp", "requests", "tqdm", "numpy", "IPython", "setproctitle"] + +[project.optional-dependencies] +runtime_common = [ + "blobfile==3.0.0", + "build", + "compressed-tensors", + "datasets", + "einops", + "fastapi", + "hf_transfer", + "huggingface_hub", + "interegular", + "llguidance>=0.7.11,<0.8.0", + "modelscope", + "msgspec", + "ninja", + "openai==1.99.1", + "openai-harmony==0.0.4", + "orjson", + "outlines==0.1.11", + "packaging", + "partial_json_parser", + "pillow", + "prometheus-client>=0.20.0", + "psutil", + "pybase64", + "pydantic", + "pynvml", + "python-multipart", + "pyzmq>=25.1.2", + "sentencepiece", + "soundfile==0.13.1", + "scipy", + "timm==1.0.16", + "tiktoken", + "torchao==0.9.0", + "transformers==4.56.1", + "uvicorn", + "uvloop", + "xgrammar==0.1.24", +] + +srt = [ + "sglang[runtime_common]", + "sgl-kernel==0.3.9.post2", + "torch==2.8.0", + "torchaudio==2.8.0", + "torchvision", + "cuda-python", + "flashinfer_python==0.3.1", +] + +blackwell = [ + "sglang[runtime_common]", + "sgl-kernel==0.3.9.post2", + "torch==2.8.0", + "torchaudio==2.8.0", + "torchvision", + "cuda-python", + "flashinfer_python==0.3.1", + "nvidia-cutlass-dsl==4.1.0", +] + +# HIP (Heterogeneous-computing Interface for Portability) for AMD +# => base docker rocm/vllm-dev:20250114, not from public vllm whl +srt_hip = [ + "sglang[runtime_common]", + "torch", + "petit_kernel==0.0.2", + "wave-lang==3.7.0", +] + +# https://docs.sglang.ai/platforms/cpu_server.html +srt_cpu = ["sglang[runtime_common]", "intel-openmp"] + +# https://docs.sglang.ai/platforms/ascend_npu.html +srt_npu = ["sglang[runtime_common]"] + +# xpu is not enabled in public vllm and torch whl, +# need to follow https://docs.vllm.ai/en/latest/getting_started/xpu-installation.htmlinstall vllm +srt_xpu = ["sglang[runtime_common]"] + +# For Intel Gaudi(device : hpu) follow the installation guide +# https://docs.vllm.ai/en/latest/getting_started/gaudi-installation.html +srt_hpu = ["sglang[runtime_common]"] + +openai = ["openai==1.99.1", "tiktoken"] +anthropic = ["anthropic>=0.20.0"] +litellm = ["litellm>=1.0.0"] +torch_memory_saver = ["torch_memory_saver==0.0.8"] +decord = ["decord"] +test = [ + "accelerate", + "expecttest", + "jsonlines", + "matplotlib", + "pandas", + "peft", + "sentence_transformers", + "pytest", + "tabulate", +] +all = ["sglang[srt]", "sglang[openai]", "sglang[anthropic]", "sglang[torch_memory_saver]", "sglang[decord]"] +all_hip = ["sglang[srt_hip]", "sglang[openai]", "sglang[anthropic]", "sglang[decord]"] +all_xpu = ["sglang[srt_xpu]", "sglang[openai]", "sglang[anthropic]", "sglang[decord]"] +all_hpu = ["sglang[srt_hpu]", "sglang[openai]", "sglang[anthropic]", "sglang[decord]"] +all_cpu = ["sglang[srt_cpu]", "sglang[openai]", "sglang[anthropic]", "sglang[decord]"] +all_npu = ["sglang[srt_npu]", "sglang[openai]", "sglang[anthropic]", "sglang[decord]"] + +dev = ["sglang[all]", "sglang[test]"] +dev_hip = ["sglang[all_hip]", "sglang[test]"] +dev_xpu = ["sglang[all_xpu]", "sglang[test]"] +dev_hpu = ["sglang[all_hpu]", "sglang[test]"] +dev_cpu = ["sglang[all_cpu]", "sglang[test]"] + +[project.urls] +"Homepage" = "https://github.com/sgl-project/sglang" +"Bug Tracker" = "https://github.com/sgl-project/sglang/issues" + +[tool.setuptools.package-data] +"sglang" = [ + "srt/layers/moe/fused_moe_triton/configs/*/*.json", + "srt/layers/quantization/configs/*.json", + "srt/mem_cache/storage/hf3fs/hf3fs_utils.cpp", +] + +[tool.setuptools.packages.find] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.wheel] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.codespell] +ignore-words-list = "ans, als, hel, boostrap, childs, te, vas, hsa, ment" +skip = "*.json,*.jsonl,*.patch,*.txt" diff --git a/python/sglang/README.md b/python/sglang/README.md new file mode 100644 index 000000000..ae0c479b9 --- /dev/null +++ b/python/sglang/README.md @@ -0,0 +1,16 @@ +# Code Structures + +- `eval`: The evaluation utilities. +- `lang`: The frontend language. +- `srt`: The backend engine for running local models. (SRT = SGLang Runtime). +- `test`: The test utilities. +- `api.py`: The public APIs. +- `bench_offline_throughput.py`: Benchmark the performance in the offline mode. +- `bench_one_batch.py`: Benchmark the latency of running a single static batch without a server. +- `bench_one_batch_server.py`: Benchmark the latency of running a single batch with a server. +- `bench_serving.py`: Benchmark online serving with dynamic requests. +- `check_env.py`: Check the environment variables and dependencies. +- `global_config.py`: The global configs and constants. +- `launch_server.py`: The entry point for launching the local server. +- `utils.py`: Common utilities. +- `version.py`: Version info. diff --git a/python/sglang/__init__.py b/python/sglang/__init__.py new file mode 100644 index 000000000..509b145a9 --- /dev/null +++ b/python/sglang/__init__.py @@ -0,0 +1,83 @@ +# SGLang public APIs + +# Frontend Language APIs +from sglang.global_config import global_config +from sglang.lang.api import ( + Engine, + Runtime, + assistant, + assistant_begin, + assistant_end, + flush_cache, + function, + gen, + gen_int, + gen_string, + get_server_info, + image, + select, + separate_reasoning, + set_default_backend, + system, + system_begin, + system_end, + user, + user_begin, + user_end, + video, +) +from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint +from sglang.lang.choices import ( + greedy_token_selection, + token_length_normalized, + unconditional_likelihood_normalized, +) + +# Lazy import some libraries +from sglang.utils import LazyImport +from sglang.version import __version__ + +Anthropic = LazyImport("sglang.lang.backend.anthropic", "Anthropic") +LiteLLM = LazyImport("sglang.lang.backend.litellm", "LiteLLM") +OpenAI = LazyImport("sglang.lang.backend.openai", "OpenAI") +VertexAI = LazyImport("sglang.lang.backend.vertexai", "VertexAI") + +# Runtime Engine APIs +ServerArgs = LazyImport("sglang.srt.server_args", "ServerArgs") +Engine = LazyImport("sglang.srt.entrypoints.engine", "Engine") + +__all__ = [ + "Engine", + "Runtime", + "assistant", + "assistant_begin", + "assistant_end", + "flush_cache", + "function", + "gen", + "gen_int", + "gen_string", + "get_server_info", + "image", + "select", + "separate_reasoning", + "set_default_backend", + "system", + "system_begin", + "system_end", + "user", + "user_begin", + "user_end", + "video", + "RuntimeEndpoint", + "greedy_token_selection", + "token_length_normalized", + "unconditional_likelihood_normalized", + "ServerArgs", + "Anthropic", + "LiteLLM", + "OpenAI", + "VertexAI", + "global_config", + "__version__", +] diff --git a/python/sglang/bench_offline_throughput.py b/python/sglang/bench_offline_throughput.py new file mode 100644 index 000000000..457d120d9 --- /dev/null +++ b/python/sglang/bench_offline_throughput.py @@ -0,0 +1,452 @@ +""" +Benchmark the throughput in the offline mode. +It accepts server arguments (the same as launch_server.py) and benchmark arguments (the same as bench_serving.py). + +# Usage +## Sharegpt dataset with default args +python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10 + +## Random dataset with default args +python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 +""" + +import argparse +import asyncio +import dataclasses +import inspect +import json +import logging +import os +import random +import time +from typing import Dict, List, Optional + +import numpy as np + +from sglang.bench_serving import ( + DatasetRow, + get_dataset, + get_tokenizer, + sample_random_requests, + set_ulimit, +) +from sglang.lang.backend.runtime_endpoint import Runtime +from sglang.srt.entrypoints.engine import Engine +from sglang.srt.server_args import ServerArgs + + +@dataclasses.dataclass +class BenchArgs: + backend: str = "engine" + result_filename: str = "" + dataset_name: str = "sharegpt" + dataset_path: str = "" + num_prompts: int = 1000 + sharegpt_output_len: Optional[int] = None + sharegpt_context_len: Optional[int] = None + random_input_len: int = 1024 + random_output_len: int = 1024 + random_range_ratio: float = 0.0 + gsp_num_groups: int = 64 + gsp_prompts_per_group: int = 16 + gsp_system_prompt_len: int = 2048 + gsp_question_len: int = 128 + gsp_output_len: int = 256 + seed: int = 1 + disable_ignore_eos: bool = False + extra_request_body: Optional[str] = None + apply_chat_template: bool = False + profile: bool = False + skip_warmup: bool = False + do_not_exit: bool = False + prompt_suffix: str = "" + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--backend", type=str, default=BenchArgs.backend) + parser.add_argument( + "--result-filename", type=str, default=BenchArgs.result_filename + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=["sharegpt", "random", "generated-shared-prefix"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument( + "--dataset-path", type=str, default="", help="Path to the dataset." + ) + parser.add_argument( + "--num-prompts", + type=int, + default=BenchArgs.num_prompts, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--sharegpt-output-len", + type=int, + default=BenchArgs.sharegpt_output_len, + help="Output length for each request. Overrides the output length from the ShareGPT dataset.", + ) + parser.add_argument( + "--sharegpt-context-len", + type=int, + default=BenchArgs.sharegpt_context_len, + help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=BenchArgs.random_input_len, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + type=int, + default=BenchArgs.random_output_len, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=BenchArgs.random_range_ratio, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + parser.add_argument( + "--gsp-num-groups", + type=int, + default=BenchArgs.gsp_num_groups, + help="Number of groups with shared prefix, used" + "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-prompts-per-group", + type=int, + default=BenchArgs.gsp_prompts_per_group, + help="Number of prompts per group of shared prefix, used" + "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-system-prompt-len", + type=int, + default=BenchArgs.gsp_system_prompt_len, + help="System prompt length, used" "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-question-len", + type=int, + default=BenchArgs.gsp_question_len, + help="Question length, used" "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-output-len", + type=int, + default=BenchArgs.gsp_output_len, + help="Target length in tokens for outputs in generated-shared-prefix dataset", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--disable-ignore-eos", + action="store_true", + help="Disable ignore EOS token", + ) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + default=BenchArgs.extra_request_body, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + parser.add_argument( + "--apply-chat-template", + action="store_true", + help="Apply chat template", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "SGLANG_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--skip-warmup", + action="store_true", + help="Skip the warmup batches.", + ) + parser.add_argument( + "--do-not-exit", + action="store_true", + help="Do not exit the program. This is useful for nsys profile with --duration and --delay.", + ) + parser.add_argument( + "--prompt-suffix", + type=str, + default="", + help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.", + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + attrs = [attr.name for attr in dataclasses.fields(cls)] + return cls(**{attr: getattr(args, attr) for attr in attrs}) + + +def throughput_test_once( + backend_name: str, + backend, + reqs: List[DatasetRow], + ignore_eos: bool, + extra_request_body: Dict, + profile: bool, +): + measurement_results = { + "backend": backend_name, + "successful_requests": len(reqs), + "total_latency": -1, + "total_input_tokens": sum(r.prompt_len for r in reqs), + "total_output_tokens": -1, + "request_throughput": -1, + "input_throughput": -1, + "output_throughput": -1, + "total_throughput": -1, + } + + prompt = [r.prompt for r in reqs] + sampling_params = [ + { + "temperature": 0, + "max_new_tokens": r.output_len, + "ignore_eos": ignore_eos, + **extra_request_body, + } + for r in reqs + ] + + if profile: + assert ( + "SGLANG_TORCH_PROFILER_DIR" in os.environ + ), "Please set SGLANG_TORCH_PROFILER_DIR." + os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True) + backend.start_profile() + + st = time.perf_counter() + gen_out = backend.generate(prompt=prompt, sampling_params=sampling_params) + latency = time.perf_counter() - st + + if profile: + dir = os.getenv("SGLANG_TORCH_PROFILER_DIR") + known_files = set(os.listdir(dir)) + backend.stop_profile() + monitor_trace_file(known_files, dir) + + if backend_name == "runtime": + gen_out = json.loads(gen_out) + + server_info = backend.get_server_info() + + measurement_results["total_latency"] = latency + measurement_results["total_output_tokens"] = sum( + o["meta_info"]["completion_tokens"] for o in gen_out + ) + measurement_results["request_throughput"] = ( + measurement_results["successful_requests"] / latency + ) + measurement_results["input_throughput"] = ( + measurement_results["total_input_tokens"] / latency + ) + measurement_results["output_throughput"] = ( + measurement_results["total_output_tokens"] / latency + ) + measurement_results["total_throughput"] = ( + measurement_results["total_input_tokens"] + + measurement_results["total_output_tokens"] + ) / latency + + if inspect.isawaitable(server_info): + server_info = asyncio.run(server_info) + + measurement_results["last_gen_throughput"] = server_info["internal_states"][0][ + "last_gen_throughput" + ] + + return measurement_results + + +def monitor_trace_file(known_files, directory, interval=1): + print(f"Monitoring {directory} for new trace files...") + + while True: + flag = False + time.sleep(interval) + current_files = set(os.listdir(directory)) + + new_files = current_files - known_files + for new_file in new_files: + new_file_path = os.path.join(directory, new_file) + print(f"New file detected: {new_file}") + + previous_size = 0 + while True: + try: + current_size = os.path.getsize(new_file_path) + except FileNotFoundError: + print(f"File {new_file} is no longer accessible.") + break + + if current_size > previous_size: + previous_size = current_size + else: + flag = True + break + + time.sleep(interval) + if flag: + break + + +def throughput_test( + server_args: ServerArgs, + bench_args: BenchArgs, +): + if bench_args.backend == "engine": + backend = Engine(**dataclasses.asdict(server_args)) + if not backend: + raise ValueError("Please provide valid engine arguments") + elif bench_args.backend == "runtime": + backend = Runtime(**dataclasses.asdict(server_args)) + else: + raise ValueError('Please set backend to either "engine" or "runtime"') + + tokenizer_id = server_args.tokenizer_path or server_args.model_path + tokenizer = get_tokenizer(tokenizer_id) + + # Set global environments + set_ulimit() + random.seed(bench_args.seed) + np.random.seed(bench_args.seed) + + # Parse args + extra_request_body = {} + if bench_args.extra_request_body: + extra_request_body = json.loads(args.extra_request_body) + + # Read dataset + input_requests = get_dataset(bench_args, tokenizer) + + warmup_requests = sample_random_requests( + input_len=256, + output_len=16, + num_prompts=min(bench_args.num_prompts, 16), + range_ratio=1.0, + tokenizer=tokenizer, + dataset_path=bench_args.dataset_path, + ) + + # Warm up + if not bench_args.skip_warmup: + logging.info("\nWarmup...") + throughput_test_once( + backend_name=bench_args.backend, + backend=backend, + reqs=warmup_requests, + ignore_eos=not bench_args.disable_ignore_eos, + extra_request_body=extra_request_body, + profile=False, + ) + time.sleep(0.5) + + logging.info("\nBenchmark...") + result = throughput_test_once( + backend_name=bench_args.backend, + backend=backend, + reqs=input_requests, + ignore_eos=not bench_args.disable_ignore_eos, + extra_request_body=extra_request_body, + profile=bench_args.profile, + ) + backend.shutdown() + + if bench_args.result_filename: + with open(bench_args.result_filename, "a") as fout: + fout.write(json.dumps(result) + "\n") + + print( + "\n{s:{c}^{n}}".format(s=" Offline Throughput Benchmark Result ", n=50, c="=") + ) + print("{:<40} {:<10}".format("Backend:", result["backend"])) + print("{:<40} {:<10}".format("Successful requests:", result["successful_requests"])) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", result["total_latency"])) + print("{:<40} {:<10}".format("Total input tokens:", result["total_input_tokens"])) + print( + "{:<40} {:<10}".format("Total generated tokens:", result["total_output_tokens"]) + ) + print( + "{:<40} {:<10.2f}".format( + "Last generation throughput (tok/s):", result["last_gen_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", result["request_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", result["input_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", result["output_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Total token throughput (tok/s):", result["total_throughput"] + ) + ) + print("=" * 50) + + return result + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + + # handling ModelScope model downloads + if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() in ("true", "1"): + if os.path.exists(args.model_path): + print(f"Using local model path: {args.model_path}") + else: + try: + from modelscope import snapshot_download + + print(f"Using ModelScope to download model: {args.model_path}") + + # download the model and replace args.model_path + args.model_path = snapshot_download( + args.model_path, + ) + print(f"Model downloaded to: {args.model_path}") + except Exception as e: + print(f"ModelScope download failed: {str(e)}") + raise e + + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + logging.basicConfig( + level=getattr(logging, server_args.log_level.upper()), + format="%(message)s", + ) + + throughput_test(server_args, bench_args) + + while bench_args.do_not_exit: + pass diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py new file mode 100644 index 000000000..ebd461ec3 --- /dev/null +++ b/python/sglang/bench_one_batch.py @@ -0,0 +1,665 @@ +""" +Benchmark the latency of running a single static batch without a server. + +This script does not launch a server and uses the low-level APIs. +It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths). + +# Usage (latency test) +## with dummy weights: +python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --load-format dummy +## sweep through multiple data points and store (append) the results in a jsonl file: +python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --output-len 32 256 --run-name test_run +## run with profiling: +python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --profile +# Usage (correctness test): +python -m sglang.bench_one_batch --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --correct + +## Reference output (of the correctness test above, can be gpu dependent): +input_ids=[[1, 450, 7483, 310, 3444, 338], [1, 450, 7483, 310, 278, 3303, 13187, 290, 338], [1, 20628, 338, 263, 6575, 1460, 2462, 322, 306, 763]] + +prefill logits (first half): tensor([[-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633], + [-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633], + [ -9.1875, -10.2500, 2.7129, ..., -4.3359, -4.0664, -4.1328]], + device='cuda:0') + +prefill logits (final): tensor([[-8.3125, -7.1172, 3.3457, ..., -4.9570, -4.1328, -3.4141], + [-8.9141, -9.0156, 4.1445, ..., -4.9922, -4.4961, -4.0781], + [-9.6328, -9.0547, 4.0195, ..., -5.3047, -4.7148, -4.4570]], + device='cuda:0') + +========== Prompt 0 ========== + The capital of France is Paris. +The capital of the United States is Washington, D.C. + + +========== Prompt 1 ========== + The capital of the United Kindom is London. +The capital of the United Kingdom is London. +The capital of the + +========== Prompt 2 ========== + Today is a sunny day and I like to go for a walk in the park. +I'm going to the park +""" + +import argparse +import copy +import dataclasses +import itertools +import json +import logging +import multiprocessing +import os +import time +from typing import Tuple + +import numpy as np +import torch +import torch.distributed as dist + +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.distributed.parallel_state import destroy_distributed_environment +from sglang.srt.entrypoints.engine import _set_envs_and_config +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.srt.layers.moe import initialize_moe_config +from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.utils import ( + configure_logger, + get_bool_env_var, + kill_process_tree, + require_mlp_sync, + require_mlp_tp_gather, + set_gpu_proc_affinity, + suppress_other_loggers, +) + + +@dataclasses.dataclass +class BenchArgs: + run_name: str = "default" + batch_size: Tuple[int] = (1,) + input_len: Tuple[int] = (1024,) + output_len: Tuple[int] = (16,) + prompt_filename: str = "" + result_filename: str = "result.jsonl" + correctness_test: bool = False + # This is only used for correctness test + cut_len: int = 4 + log_decode_step: int = 0 + profile: bool = False + profile_record_shapes: bool = False + profile_filename_prefix: str = "profile" + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--run-name", type=str, default=BenchArgs.run_name) + parser.add_argument( + "--batch-size", type=int, nargs="+", default=BenchArgs.batch_size + ) + parser.add_argument( + "--input-len", type=int, nargs="+", default=BenchArgs.input_len + ) + parser.add_argument( + "--output-len", type=int, nargs="+", default=BenchArgs.output_len + ) + parser.add_argument( + "--prompt-filename", type=str, default=BenchArgs.prompt_filename + ) + parser.add_argument( + "--result-filename", type=str, default=BenchArgs.result_filename + ) + parser.add_argument("--correctness-test", action="store_true") + parser.add_argument("--cut-len", type=int, default=BenchArgs.cut_len) + parser.add_argument( + "--log-decode-step", + type=int, + default=BenchArgs.log_decode_step, + help="Log decode latency by step, default is set to zero to disable.", + ) + parser.add_argument( + "--profile", action="store_true", help="Use Torch Profiler." + ) + parser.add_argument( + "--profile-record-shapes", + action="store_true", + help="Record tensor shapes in profiling results.", + ) + parser.add_argument( + "--profile-filename-prefix", + type=str, + default=BenchArgs.profile_filename_prefix, + help="Prefix of the profiling file names. The full profiling result file(s) be " + '"[profile_filename_prefix]_batch[batch_size]_input[input_len]_output[output_len].trace.json.gz"', + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + # use the default value's type to cast the args into correct types. + attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] + return cls( + **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} + ) + + +def load_model(server_args, port_args, tp_rank): + suppress_other_loggers() + rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None + moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size) + + model_config = ModelConfig.from_server_args(server_args) + model_runner = ModelRunner( + model_config=model_config, + mem_fraction_static=server_args.mem_fraction_static, + gpu_id=tp_rank, + tp_rank=tp_rank, + tp_size=server_args.tp_size, + moe_ep_rank=moe_ep_rank, + moe_ep_size=server_args.ep_size, + pp_rank=0, + pp_size=1, + nccl_port=port_args.nccl_port, + server_args=server_args, + ) + rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}") + tokenizer = get_tokenizer( + server_args.tokenizer_path, + tokenizer_mode=server_args.tokenizer_mode, + trust_remote_code=server_args.trust_remote_code, + ) + if server_args.tp_size > 1: + dist.barrier() + return model_runner, tokenizer + + +def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts): + prompts = ( + custom_prompts + if custom_prompts + else [ + "The capital of France is", + "The capital of the United Kindom is", + "Today is a sunny day and I like", + ] + ) + input_ids = [tokenizer.encode(p) for p in prompts] + sampling_params = SamplingParams( + temperature=0, + max_new_tokens=BenchArgs.output_len, + ) + + reqs = [] + for i in range(len(prompts)): + assert len(input_ids[i]) > bench_args.cut_len + + tmp_input_ids = input_ids[i][: bench_args.cut_len] + req = Req( + rid=i, + origin_input_text=prompts[i], + origin_input_ids=tmp_input_ids, + sampling_params=sampling_params, + ) + req.prefix_indices = [] + req.fill_ids = req.origin_input_ids + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + reqs.append(req) + + return input_ids, reqs + + +def prepare_extend_inputs_for_correctness_test( + bench_args, input_ids, reqs, model_runner +): + for i in range(len(reqs)): + req = reqs[i] + req.fill_ids += input_ids[i][bench_args.cut_len :] + req.prefix_indices = model_runner.req_to_token_pool.req_to_token[ + i, : bench_args.cut_len + ] + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + return reqs + + +def prepare_synthetic_inputs_for_latency_test( + batch_size, input_len, custom_inputs=None +): + input_ids = ( + custom_inputs + if custom_inputs + else np.random.randint(0, 10000, (batch_size, input_len), dtype=np.int32) + ) + sampling_params = SamplingParams( + temperature=0, + max_new_tokens=BenchArgs.output_len, + ) + + reqs = [] + for i in range(len(input_ids)): + req = Req( + rid=i, + origin_input_text="", + origin_input_ids=list(input_ids[i]), + sampling_params=sampling_params, + ) + req.prefix_indices = [] + req.fill_ids = req.origin_input_ids + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + reqs.append(req) + + return reqs + + +@torch.no_grad +def extend(reqs, model_runner): + batch = ScheduleBatch.init_new( + reqs=reqs, + req_to_token_pool=model_runner.req_to_token_pool, + token_to_kv_pool_allocator=model_runner.token_to_kv_pool_allocator, + tree_cache=None, + model_config=model_runner.model_config, + enable_overlap=False, + spec_algorithm=SpeculativeAlgorithm.NONE, + ) + batch.prepare_for_extend() + _maybe_prepare_mlp_sync_batch(batch, model_runner) + model_worker_batch = batch.get_model_worker_batch() + forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner) + logits_output, _ = model_runner.forward(forward_batch) + next_token_ids = model_runner.sample(logits_output, forward_batch) + return next_token_ids, logits_output.next_token_logits, batch + + +@torch.no_grad +def decode(input_token_ids, batch, model_runner): + batch.output_ids = input_token_ids + batch.prepare_for_decode() + _maybe_prepare_mlp_sync_batch(batch, model_runner) + model_worker_batch = batch.get_model_worker_batch() + forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner) + logits_output, _ = model_runner.forward(forward_batch) + next_token_ids = model_runner.sample(logits_output, forward_batch) + return next_token_ids, logits_output.next_token_logits + + +def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner): + if require_mlp_sync(model_runner.server_args): + Scheduler.prepare_mlp_sync_batch_raw( + batch, + dp_size=model_runner.server_args.dp_size, + attn_tp_size=1, + tp_group=model_runner.tp_group, + get_idle_batch=None, + disable_cuda_graph=model_runner.server_args.disable_cuda_graph, + spec_algorithm=SpeculativeAlgorithm.NONE, + speculative_num_draft_tokens=None, + require_mlp_tp_gather=require_mlp_tp_gather(model_runner.server_args), + disable_overlap_schedule=model_runner.server_args.disable_overlap_schedule, + ) + + +def _read_prompts_from_file(prompt_file, rank_print): + """Read custom prompts from the file specified by `--prompt-filename`.""" + if not prompt_file: + return [] + if not os.path.exists(prompt_file): + rank_print( + f"Custom prompt file {prompt_file} not found. Using default inputs..." + ) + return [] + with open(prompt_file, "r") as pf: + return pf.readlines() + + +def _save_profile_trace_results(profiler, filename): + parent_dir = os.path.dirname(os.path.abspath(filename)) + os.makedirs(parent_dir, exist_ok=True) + profiler.export_chrome_trace(filename) + print( + profiler.key_averages(group_by_input_shape=True).table( + sort_by="self_cpu_time_total" + ) + ) + + +def correctness_test( + server_args, + port_args, + bench_args, + tp_rank, +): + # Configure the logger + configure_logger(server_args, prefix=f" TP{tp_rank}") + rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None + + # Load the model + model_runner, tokenizer = load_model(server_args, port_args, tp_rank) + + # Prepare inputs + custom_prompts = _read_prompts_from_file(bench_args.prompt_filename, rank_print) + input_ids, reqs = prepare_inputs_for_correctness_test( + bench_args, tokenizer, custom_prompts + ) + rank_print(f"\n{input_ids=}\n") + + if bench_args.cut_len > 0: + # Prefill + next_token_ids, next_token_logits, batch = extend(reqs, model_runner) + rank_print(f"prefill logits (first half): {next_token_logits} \n") + + # Prepare extend inputs + reqs = prepare_extend_inputs_for_correctness_test( + bench_args, input_ids, reqs, model_runner + ) + + # Extend (prefill w/ KV cache) + next_token_ids, next_token_logits, batch = extend(reqs, model_runner) + rank_print(f"prefill logits (final): {next_token_logits} \n") + + # Decode + output_ids = [input_ids[i] + [next_token_ids[i]] for i in range(len(input_ids))] + for _ in range(bench_args.output_len[0] - 1): + next_token_ids, _ = decode(next_token_ids, batch, model_runner) + next_token_ids_list = next_token_ids.tolist() + for i in range(len(reqs)): + output_ids[i].append(next_token_ids_list[i]) + + # Print output texts + for i in range(len(reqs)): + rank_print(f"========== Prompt {i} ==========") + rank_print(tokenizer.decode(output_ids[i]), "\n") + + +def synchronize(device): + torch.get_device_module(device).synchronize() + + +def latency_test_run_once( + run_name, + model_runner, + rank_print, + reqs, + batch_size, + input_len, + output_len, + device, + log_decode_step, + profile, + profile_record_shapes, + profile_filename_prefix, +): + max_batch_size = model_runner.max_total_num_tokens // (input_len + output_len) + if batch_size > max_batch_size: + rank_print( + f"skipping ({batch_size}, {input_len}, {output_len}) due to max batch size limit" + ) + return + + # Clear the pools. + model_runner.req_to_token_pool.clear() + model_runner.token_to_kv_pool_allocator.clear() + + measurement_results = { + "run_name": run_name, + "batch_size": batch_size, + "input_len": input_len, + "output_len": output_len, + } + + tot_latency = 0 + + profiler = None + if profile: + profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, + record_shapes=profile_record_shapes, + ) + profiler.start() + + # Prefill + synchronize(device) + tic = time.perf_counter() + next_token_ids, _, batch = extend(reqs, model_runner) + synchronize(device) + prefill_latency = time.perf_counter() - tic + tot_latency += prefill_latency + throughput = input_len * batch_size / prefill_latency + rank_print( + f"Prefill. latency: {prefill_latency:6.5f} s, throughput: {throughput:9.2f} token/s" + ) + measurement_results["prefill_latency"] = prefill_latency + measurement_results["prefill_throughput"] = throughput + + if profile: + profiler.stop() + profile_filename = f"{profile_filename_prefix}_batch{batch_size}_input{input_len}_output{output_len}_prefill.trace.json.gz" + _save_profile_trace_results(profiler, profile_filename) + rank_print( + f"torch profiler chrome trace for prefill saved to {profile_filename}" + ) + + # Decode + decode_latencies = [] + for i in range(output_len - 1): + synchronize(device) + if profile and i == output_len / 2: + profiler = None + profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, + record_shapes=profile_record_shapes, + ) + profiler.start() + + tic = time.perf_counter() + next_token_ids, _ = decode(next_token_ids, batch, model_runner) + synchronize(device) + latency = time.perf_counter() - tic + tot_latency += latency + throughput = batch_size / latency + decode_latencies.append(latency) + if i < 5 or (log_decode_step > 0 and i % log_decode_step == 0): + rank_print( + f"Decode {i}. Batch size: {batch_size}, latency: {latency:6.5f} s, throughput: {throughput:9.2f} token/s" + ) + + if profile and i == output_len / 2: + profiler.stop() + profile_filename = f"{profile_filename_prefix}_batch{batch_size}_input{input_len}_output{output_len}_decode.trace.json.gz" + _save_profile_trace_results(profiler, profile_filename) + rank_print( + f"torch profiler chrome trace for decoding 1 token saved to {profile_filename}" + ) + + # Record decode timing from 2nd output + if output_len > 1: + med_decode_latency = np.median(decode_latencies) + med_decode_throughput = batch_size / med_decode_latency + rank_print( + f"Decode. median latency: {med_decode_latency:6.5f} s, median throughput: {med_decode_throughput:9.2f} token/s" + ) + measurement_results["median_decode_latency"] = med_decode_latency + measurement_results["median_decode_throughput"] = med_decode_throughput + + throughput = (input_len + output_len) * batch_size / tot_latency + rank_print( + f"Total. latency: {tot_latency:6.3f} s, throughput: {throughput:9.2f} token/s" + ) + measurement_results["total_latency"] = tot_latency + measurement_results["overall_throughput"] = throughput + return measurement_results + + +def latency_test( + server_args, + port_args, + bench_args, + tp_rank, +): + initialize_moe_config(server_args) + + # Set CPU affinity + if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"): + set_gpu_proc_affinity(server_args.tp_size, server_args.nnodes, tp_rank) + + # Configure the logger + configure_logger(server_args, prefix=f" TP{tp_rank}") + rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None + + # Load the model + model_runner, tokenizer = load_model(server_args, port_args, tp_rank) + + # Prepare inputs for warm up + reqs = prepare_synthetic_inputs_for_latency_test( + bench_args.batch_size[0], bench_args.input_len[0] + ) + + # Warm up + rank_print("Warmup ...") + latency_test_run_once( + bench_args.run_name, + model_runner, + rank_print, + reqs, + bench_args.batch_size[0], + bench_args.input_len[0], + min(32, bench_args.output_len[0]), # shorter decoding to speed up the warmup + server_args.device, + log_decode_step=0, + profile=False, + profile_record_shapes=False, + profile_filename_prefix="", # not used + ) + + rank_print("Benchmark ...") + + custom_inputs = _read_prompts_from_file(bench_args.prompt_filename, rank_print) + custom_inputs = [tokenizer.encode(p.strip()) for p in custom_inputs] + custom_input_len = len(custom_inputs) + + # Run the sweep + result_list = [] + for bs, il, ol in itertools.product( + bench_args.batch_size, bench_args.input_len, bench_args.output_len + ): + bs_aligned_inputs = [] + if custom_inputs: + if custom_input_len == bs: + bs_aligned_inputs = custom_inputs + elif custom_input_len > bs: + rank_print( + f"Custom input size ({custom_input_len}) is larger than batch_size ({bs}). " + f"Using the first {bs} prompts." + ) + bs_aligned_inputs = copy.deepcopy(custom_inputs[:bs]) + else: + rank_print( + f"Custom input size ({custom_input_len}) is smaller than batch_size ({bs}). " + f"Pad to the desired batch_size with the last prompt." + ) + bs_aligned_inputs = copy.deepcopy(custom_inputs) + bs_aligned_inputs.extend( + [bs_aligned_inputs[-1]] * (bs - custom_input_len) + ) + + reqs = prepare_synthetic_inputs_for_latency_test(bs, il, bs_aligned_inputs) + ret = latency_test_run_once( + bench_args.run_name, + model_runner, + rank_print, + reqs, + bs, + il, + ol, + server_args.device, + bench_args.log_decode_step, + bench_args.profile if tp_rank == 0 else None, + bench_args.profile_record_shapes if tp_rank == 0 else None, + bench_args.profile_filename_prefix, + ) + if ret is not None: + result_list.append(ret) + + # Write results in jsonlines format on rank 0. + if tp_rank == 0 and bench_args.result_filename: + with open(bench_args.result_filename, "a") as fout: + for result in result_list: + fout.write(json.dumps(result) + "\n") + + if server_args.tp_size > 1: + destroy_distributed_environment() + + +def main(server_args, bench_args): + server_args.cuda_graph_max_bs = max(bench_args.batch_size) + + _set_envs_and_config(server_args) + + if server_args.model_path: + if bench_args.correctness_test: + work_func = correctness_test + else: + work_func = latency_test + else: + raise ValueError( + "Provide --model-path for running the tests or " + "provide --result-filename for plotting the results" + ) + + port_args = PortArgs.init_new(server_args) + + if server_args.tp_size == 1: + work_func(server_args, port_args, bench_args, 0) + else: + workers = [] + for tp_rank in range(server_args.tp_size): + proc = multiprocessing.Process( + target=work_func, + args=( + server_args, + port_args, + bench_args, + tp_rank, + ), + ) + proc.start() + workers.append(proc) + + for proc in workers: + proc.join() + + proc.terminate() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + logging.basicConfig( + level=getattr(logging, server_args.log_level.upper()), + format="%(message)s", + ) + + try: + main(server_args, bench_args) + finally: + if server_args.tp_size != 1: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/python/sglang/bench_one_batch_server.py b/python/sglang/bench_one_batch_server.py new file mode 100644 index 000000000..8495c110e --- /dev/null +++ b/python/sglang/bench_one_batch_server.py @@ -0,0 +1,439 @@ +""" +Benchmark the latency of running a single batch with a server. + +This script launches a server and uses the HTTP interface. +It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths). + +Usage: +python3 -m sglang.bench_one_batch_server --model meta-llama/Meta-Llama-3.1-8B --batch-size 1 16 64 --input-len 1024 --output-len 8 + +python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8 +python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8 --show-report --profile --profile-by-stage +""" + +import argparse +import dataclasses +import itertools +import json +import multiprocessing +import os +import time +from typing import List, Tuple + +import requests + +from sglang.bench_serving import get_tokenizer, sample_random_requests +from sglang.profiler import run_profile +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import is_blackwell, kill_process_tree +from sglang.test.test_utils import is_in_ci, write_github_step_summary + + +@dataclasses.dataclass +class BenchArgs: + run_name: str = "default" + batch_size: Tuple[int] = (1,) + input_len: Tuple[int] = (1024,) + output_len: Tuple[int] = (16,) + temperature: float = 0.0 + return_logprob: bool = False + client_stream_interval: int = 1 + input_len_step_percentage: float = 0.0 + result_filename: str = "result.jsonl" + base_url: str = "" + skip_warmup: bool = False + show_report: bool = False + profile: bool = False + profile_steps: int = 3 + profile_by_stage: bool = False + dataset_path: str = "" + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--run-name", type=str, default=BenchArgs.run_name) + parser.add_argument( + "--batch-size", type=int, nargs="+", default=BenchArgs.batch_size + ) + parser.add_argument( + "--input-len", type=int, nargs="+", default=BenchArgs.input_len + ) + parser.add_argument( + "--output-len", type=int, nargs="+", default=BenchArgs.output_len + ) + parser.add_argument("--temperature", type=float, default=BenchArgs.temperature) + parser.add_argument("--return-logprob", action="store_true") + parser.add_argument( + "--client-stream-interval", + type=int, + default=BenchArgs.client_stream_interval, + ) + parser.add_argument( + "--input-len-step-percentage", + type=float, + default=BenchArgs.input_len_step_percentage, + ) + parser.add_argument( + "--result-filename", type=str, default=BenchArgs.result_filename + ) + parser.add_argument("--base-url", type=str, default=BenchArgs.base_url) + parser.add_argument("--skip-warmup", action="store_true") + parser.add_argument("--show-report", action="store_true") + parser.add_argument("--profile", action="store_true") + parser.add_argument( + "--profile-steps", type=int, default=BenchArgs.profile_steps + ) + parser.add_argument("--profile-by-stage", action="store_true") + parser.add_argument( + "--dataset-path", + type=str, + default=BenchArgs.dataset_path, + help="Path to the dataset.", + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + # use the default value's type to cast the args into correct types. + attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] + return cls( + **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} + ) + + +def launch_server_internal(server_args): + try: + launch_server(server_args) + except Exception as e: + raise e + finally: + kill_process_tree(os.getpid(), include_parent=False) + + +def launch_server_process(server_args: ServerArgs): + proc = multiprocessing.Process(target=launch_server_internal, args=(server_args,)) + proc.start() + base_url = f"http://{server_args.host}:{server_args.port}" + timeout = 600 + + start_time = time.time() + while time.time() - start_time < timeout: + try: + headers = { + "Content-Type": "application/json; charset=utf-8", + } + response = requests.get(f"{base_url}/v1/models", headers=headers) + if response.status_code == 200: + return proc, base_url + except requests.RequestException: + pass + time.sleep(10) + raise TimeoutError("Server failed to start within the timeout period.") + + +def run_one_case( + url: str, + batch_size: int, + input_len: int, + output_len: int, + temperature: float, + return_logprob: bool, + stream_interval: int, + input_len_step_percentage: float, + run_name: str, + result_filename: str, + tokenizer, + profile: bool = False, + profile_steps: int = 3, + profile_by_stage: bool = False, + dataset_path: str = "", +): + requests.post(url + "/flush_cache") + input_requests = sample_random_requests( + input_len=input_len, + output_len=output_len, + num_prompts=batch_size, + range_ratio=1.0, + tokenizer=tokenizer, + dataset_path=dataset_path, + random_sample=True, + return_text=False, + ) + + use_structured_outputs = False + if use_structured_outputs: + texts = [] + for _ in range(batch_size): + texts.append( + "Human: What is the capital city of france? can you give as many trivial information as possible about that city? answer in json.\n" + * 50 + + "Assistant:" + ) + json_schema = "$$ANY$$" + else: + json_schema = None + + profile_link = None + if profile: + profile_link: str = run_profile( + url, profile_steps, ["CPU", "GPU"], None, None, profile_by_stage + ) + + tic = time.perf_counter() + response = requests.post( + url + "/generate", + json={ + "input_ids": [req.prompt for req in input_requests], + "sampling_params": { + "temperature": temperature, + "max_new_tokens": output_len, + "ignore_eos": True, + "json_schema": json_schema, + "stream_interval": stream_interval, + }, + "return_logprob": return_logprob, + "stream": True, + }, + stream=True, + ) + + # The TTFT of the last request in the batch + ttft = 0.0 + for chunk in response.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + if "error" in data: + raise RuntimeError(f"Request has failed. {data}.") + + assert ( + data["meta_info"]["finish_reason"] is None + or data["meta_info"]["finish_reason"]["type"] == "length" + ) + if data["meta_info"]["completion_tokens"] == 1: + ttft = time.perf_counter() - tic + + latency = time.perf_counter() - tic + input_throughput = batch_size * input_len / ttft + output_throughput = batch_size * output_len / (latency - ttft) + overall_throughput = batch_size * (input_len + output_len) / latency + + server_info = requests.get(url + "/get_server_info").json() + acc_length = server_info["internal_states"][0].get("avg_spec_accept_length", None) + last_gen_throughput = server_info["internal_states"][0]["last_gen_throughput"] + + print(f"batch size: {batch_size}") + print(f"input_len: {input_len}") + print(f"output_len: {output_len}") + print(f"latency: {latency:.2f} s") + print(f"ttft: {ttft:.2f} s") + print(f"last generation throughput: {last_gen_throughput:.2f} tok/s") + print(f"input throughput: {input_throughput:.2f} tok/s") + if output_len != 1: + print(f"output throughput: {output_throughput:.2f} tok/s") + + if result_filename: + with open(result_filename, "a") as fout: + res = { + "run_name": run_name, + "batch_size": batch_size, + "input_len": input_len, + "output_len": output_len, + "latency": round(latency, 4), + "output_throughput": round(output_throughput, 2), + "overall_throughput": round(overall_throughput, 2), + "last_gen_throughput": round(last_gen_throughput, 2), + } + fout.write(json.dumps(res) + "\n") + + return ( + batch_size, + latency, + ttft, + input_throughput, + output_throughput, + overall_throughput, + last_gen_throughput, + acc_length, + profile_link if profile else None, + ) + + +def get_report_summary( + result: List[Tuple], server_args: ServerArgs, bench_args: BenchArgs +): + import tabulate + + summary = ( + f"\nInput lens: {bench_args.input_len}. Output lens: {bench_args.output_len}.\n" + ) + + headers = [ + "batch size", + "latency (s)", + "input throughput (tok/s)", + "output throughput (tok/s)", + "acc length", + "ITL (ms)", + "input cost ($/1M)", + "output cost ($/1M)", + ] + if bench_args.profile: + headers.append("profile") + rows = [] + + for ( + batch_size, + latency, + ttft, + input_throughput, + output_throughput, + _, + _, + acc_length, + trace_link, + ) in result: + if is_blackwell(): + hourly_cost_per_gpu = 4 # $4/hour for one B200 + else: + hourly_cost_per_gpu = 2 # $2/hour for one H100 + + hourly_cost = hourly_cost_per_gpu * server_args.tp_size + input_util = 0.7 + accept_length = round(acc_length, 2) if acc_length is not None else "n/a" + itl = 1 / (output_throughput / batch_size) * 1000 + input_cost = 1e6 / (input_throughput * input_util) / 3600 * hourly_cost + output_cost = 1e6 / output_throughput / 3600 * hourly_cost + row = [ + batch_size, + latency, + input_throughput, + output_throughput, + accept_length, + itl, + input_cost, + output_cost, + ] + if trace_link: + row.append(f"[Profile]({trace_link})") + rows.append(row) + + summary += tabulate.tabulate( + rows, headers=headers, tablefmt="github", floatfmt=".2f" + ) + return summary + + +def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs): + if bench_args.base_url: + proc, base_url = None, bench_args.base_url + else: + proc, base_url = launch_server_process(server_args) + + server_info = requests.get(base_url + "/get_server_info").json() + if "tokenizer_path" in server_info: + tokenizer_path = server_info["tokenizer_path"] + elif "prefill" in server_info: + tokenizer_path = server_info["prefill"][0]["tokenizer_path"] + tokenizer = get_tokenizer(tokenizer_path) + + # warmup + if not bench_args.skip_warmup: + print("=" * 8 + " Warmup Begin " + "=" * 8) + run_one_case( + base_url, + batch_size=16, + input_len=1024, + output_len=16, + temperature=bench_args.temperature, + return_logprob=bench_args.return_logprob, + stream_interval=bench_args.client_stream_interval, + input_len_step_percentage=bench_args.input_len_step_percentage, + run_name="", + result_filename="", + tokenizer=tokenizer, + dataset_path=bench_args.dataset_path, + ) + print("=" * 8 + " Warmup End " + "=" * 8 + "\n") + + # benchmark + result = [] + bench_result = [] + try: + for bs, il, ol in itertools.product( + bench_args.batch_size, bench_args.input_len, bench_args.output_len + ): + result.append( + run_one_case( + base_url, + bs, + il, + ol, + temperature=bench_args.temperature, + return_logprob=bench_args.return_logprob, + stream_interval=bench_args.client_stream_interval, + input_len_step_percentage=bench_args.input_len_step_percentage, + run_name=bench_args.run_name, + result_filename=bench_args.result_filename, + tokenizer=tokenizer, + ) + ) + + if bench_args.profile: + try: + for bs, il, ol in itertools.product( + bench_args.batch_size, bench_args.input_len, bench_args.output_len + ): + bench_result.append( + ( + run_one_case( + base_url, + bs, + il, + ol, + temperature=bench_args.temperature, + return_logprob=bench_args.return_logprob, + stream_interval=bench_args.client_stream_interval, + input_len_step_percentage=bench_args.input_len_step_percentage, + run_name=bench_args.run_name, + result_filename=bench_args.result_filename, + tokenizer=tokenizer, + profile=bench_args.profile, + profile_steps=bench_args.profile_steps, + profile_by_stage=bench_args.profile_by_stage, + )[-1], + ) + ) + result = [t1[:-1] + t2 for t1, t2 in zip(result, bench_result)] + except Exception as e: + print(f"Error profiling, there will be no profile trace dump: {e}") + finally: + if proc: + kill_process_tree(proc.pid) + + print(f"\nResults are saved to {bench_args.result_filename}") + + if not bench_args.show_report: + return + + summary = get_report_summary(result, server_args, bench_args) + print(summary) + + if is_in_ci(): + write_github_step_summary(summary) + + +def main(): + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + run_benchmark(server_args, bench_args) + + +if __name__ == "__main__": + main() diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py new file mode 100644 index 000000000..27ff8a6da --- /dev/null +++ b/python/sglang/bench_serving.py @@ -0,0 +1,2403 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/backend_request_func.py +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/benchmark_serving.py + +""" +Benchmark online serving with dynamic requests. + +Usage: +python3 -m sglang.bench_serving --backend sglang --num-prompt 10 + +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5 +""" + +import argparse +import asyncio +import base64 +import io +import json +import os +import pickle +import random +import resource +import sys +import time +import traceback +import warnings +from argparse import ArgumentParser +from dataclasses import dataclass, field +from datetime import datetime +from json import JSONDecodeError +from pathlib import Path +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union + +import aiohttp +import numpy as np +import requests +from tqdm.asyncio import tqdm +from transformers import ( + AutoTokenizer, + PreTrainedTokenizer, + PreTrainedTokenizerBase, + PreTrainedTokenizerFast, +) + +ASSISTANT_SUFFIX = "Assistant:" + +global args + + +# don't want to import sglang package here +def _get_bool_env_var(name: str, default: str = "false") -> bool: + value = os.getenv(name, default) + return value.lower() in ("true", "1") + + +def _create_bench_client_session(): + # When the pressure is big, the read buffer could be full before aio thread read + # the content. We increase the read_bufsize from 64K to 10M. + # Define constants for timeout and buffer size for clarity and maintainability + BENCH_AIOHTTP_TIMEOUT_SECONDS = 6 * 60 * 60 # 6 hours + BENCH_AIOHTTP_READ_BUFSIZE_BYTES = 10 * 1024**2 # 10 MB + + aiohttp_timeout = aiohttp.ClientTimeout(total=BENCH_AIOHTTP_TIMEOUT_SECONDS) + return aiohttp.ClientSession( + timeout=aiohttp_timeout, read_bufsize=BENCH_AIOHTTP_READ_BUFSIZE_BYTES + ) + + +@dataclass +class RequestFuncInput: + prompt: str + api_url: str + prompt_len: int + output_len: int + model: str + lora_name: str + image_data: Optional[List[str]] + extra_request_body: Dict[str, Any] + timestamp: Optional[float] = None + + +@dataclass +class RequestFuncOutput: + generated_text: str = "" + success: bool = False + latency: float = 0.0 + ttft: float = 0.0 # Time to first token + itl: List[float] = field(default_factory=list) # List of inter-token latencies + prompt_len: int = 0 + error: str = "" + output_len: int = 0 + + @staticmethod + def init_new(request_func_input: RequestFuncInput): + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + return output + + +def remove_prefix(text: str, prefix: str) -> str: + return text[len(prefix) :] if text.startswith(prefix) else text + + +def remove_suffix(text: str, suffix: str) -> str: + return text[: -len(suffix)] if text.endswith(suffix) else text + + +def get_auth_headers() -> Dict[str, str]: + api_key = os.environ.get("OPENAI_API_KEY") + if api_key: + return {"Authorization": f"Bearer {api_key}"} + else: + return {} + + +# trt llm does not support ignore_eos +# https://github.com/triton-inference-server/tensorrtllm_backend/issues/505 +async def async_request_trt_llm( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith("generate_stream") + + async with _create_bench_client_session() as session: + payload = { + "accumulate_tokens": True, + "text_input": request_func_input.prompt, + "temperature": 0.000001, + "top_p": 1.0, + "max_tokens": request_func_input.output_len, + "stream": True, + "min_length": request_func_input.output_len, + "end_id": 1048576, + **request_func_input.extra_request_body, + } + if args.disable_ignore_eos: + del payload["min_length"] + del payload["end_id"] + output = RequestFuncOutput.init_new(request_func_input) + + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data:") + + data = json.loads(chunk) + output.generated_text += data["text_output"] + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + + output.latency = most_recent_timestamp - st + output.success = True + output.output_len = request_func_input.output_len + + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +# set ignore_eos True by default +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + "completions" + ), "OpenAI Completions API URL must end with 'completions'." + + prompt = request_func_input.prompt + + async with _create_bench_client_session() as session: + payload = { + "model": request_func_input.model, + "prompt": prompt, + "temperature": 0.0, + "best_of": 1, + "max_tokens": request_func_input.output_len, + "stream": not args.disable_stream, + "ignore_eos": not args.disable_ignore_eos, + **request_func_input.extra_request_body, + } + headers = get_auth_headers() + + output = RequestFuncOutput.init_new(request_func_input) + + generated_text = "" + output_len = request_func_input.output_len + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["choices"][0]["text"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += data["choices"][0]["text"] + output_len = (data.get("usage") or {}).get( + "completion_tokens", output_len + ) + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_openai_chat_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + """Makes a request to the OpenAI Chat Completions API. + + Handles both streaming and non-streaming responses, including support + for image data in messages. Calculates and returns various performance + metrics. + + Args: + request_func_input: Input parameters for the request. + pbar: Optional tqdm progress bar to update. + + Returns: + RequestFuncOutput: Output of the request, including generated text, + latency, TTFT, ITL, and success status. + """ + api_url = request_func_input.api_url + assert api_url.endswith( + "chat/completions" + ), "OpenAI Chat Completions API URL must end with 'chat/completions'." + + if request_func_input.image_data: + # Build multi-image content: a list of image_url entries followed by the text + content_items = [ + { + "type": "image_url", + "image_url": {"url": img_url}, + } + for img_url in request_func_input.image_data + ] + content_items.append({"type": "text", "text": request_func_input.prompt}) + messages = [ + { + "role": "user", + "content": content_items, + }, + ] + else: + messages = [{"role": "user", "content": request_func_input.prompt}] + + async with _create_bench_client_session() as session: + payload = { + "model": request_func_input.model, + "messages": messages, + "temperature": 0.0, + "max_tokens": request_func_input.output_len, + "stream": not args.disable_stream, + **request_func_input.extra_request_body, + } + headers = get_auth_headers() + + output = RequestFuncOutput.init_new(request_func_input) + + generated_text = "" + output_len = request_func_input.output_len + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + if args.disable_stream: + # Non-streaming response + response_json = await response.json() + output.generated_text = response_json["choices"][0]["message"][ + "content" + ] + output.success = True + output.latency = time.perf_counter() - st + output.ttft = ( + output.latency + ) # For non-streaming, TTFT = total latency + output.output_len = response_json.get("usage", {}).get( + "completion_tokens", output_len + ) + else: + # Streaming response + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # Check if this chunk contains content + delta = data.get("choices", [{}])[0].get("delta", {}) + content = delta.get("content", "") + + if content: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append( + timestamp - most_recent_timestamp + ) + + most_recent_timestamp = timestamp + generated_text += content + + # Check for usage info in final chunk + output_len = (data.get("usage") or {}).get( + "completion_tokens", output_len + ) + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_truss( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + + prompt = request_func_input.prompt + + async with _create_bench_client_session() as session: + payload = { + "model": request_func_input.model, + "prompt": prompt, + "temperature": 0.0, + "best_of": 1, + "max_tokens": request_func_input.output_len, + "stream": not args.disable_stream, + "ignore_eos": not args.disable_ignore_eos, + **request_func_input.extra_request_body, + } + headers = get_auth_headers() + + output = RequestFuncOutput.init_new(request_func_input) + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["choices"][0]["text"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += data["choices"][0]["text"] + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = request_func_input.output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_sglang_generate( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + prompt = request_func_input.prompt + + async with _create_bench_client_session() as session: + payload = { + ("text" if isinstance(prompt, str) else "input_ids"): prompt, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": request_func_input.output_len, + "ignore_eos": not args.disable_ignore_eos, + }, + "stream": not args.disable_stream, + "lora_path": request_func_input.lora_name, + "return_logprob": args.return_logprob, + "logprob_start_len": -1, + **request_func_input.extra_request_body, + } + + # Add image data if available (list of image urls/base64) + if request_func_input.image_data: + payload["image_data"] = request_func_input.image_data + + headers = get_auth_headers() + + output = RequestFuncOutput.init_new(request_func_input) + + generated_text = "" + output_len = request_func_input.output_len + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + last_output_len = 0 + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if "text" in data and data["text"]: + timestamp = time.perf_counter() + generated_text = data["text"] + output_len = data["meta_info"]["completion_tokens"] + + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + num_new_tokens = output_len - last_output_len + if num_new_tokens == 0: + continue + adjust_itl = ( + timestamp - most_recent_timestamp + ) / num_new_tokens + output.itl.extend([adjust_itl] * num_new_tokens) + + most_recent_timestamp = timestamp + last_output_len = output_len + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + print(f"{output.error=}") + + if pbar: + pbar.update(1) + return output + + +async def async_request_gserver( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + raise NotImplementedError() + + +async def async_request_profile(api_url: str) -> RequestFuncOutput: + async with _create_bench_client_session() as session: + output = RequestFuncOutput() + try: + async with session.post(url=api_url) as response: + if response.status == 200: + output.success = True + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + return output + + +def get_model(pretrained_model_name_or_path: str) -> str: + if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() == "true": + import huggingface_hub.constants + from modelscope import snapshot_download + + model_path = snapshot_download( + model_id=pretrained_model_name_or_path, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"], + ) + + return model_path + return pretrained_model_name_or_path + + +def get_tokenizer( + pretrained_model_name_or_path: str, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + assert ( + pretrained_model_name_or_path is not None + and pretrained_model_name_or_path != "" + ) + if pretrained_model_name_or_path.endswith( + ".json" + ) or pretrained_model_name_or_path.endswith(".model"): + from sglang.srt.hf_transformers_utils import get_tokenizer + + return get_tokenizer(pretrained_model_name_or_path) + + if pretrained_model_name_or_path is not None and not os.path.exists( + pretrained_model_name_or_path + ): + pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) + return AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=True + ) + + +def get_dataset(args, tokenizer): + tokenize_prompt = getattr(args, "tokenize_prompt", False) + if args.dataset_name == "sharegpt": + assert not tokenize_prompt + input_requests = sample_sharegpt_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + fixed_output_len=args.sharegpt_output_len, + context_len=args.sharegpt_context_len, + prompt_suffix=args.prompt_suffix, + apply_chat_template=args.apply_chat_template, + ) + elif args.dataset_name.startswith("random") and args.dataset_name != "random-image": + input_requests = sample_random_requests( + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + dataset_path=args.dataset_path, + random_sample=args.dataset_name == "random", + return_text=not tokenize_prompt, + ) + elif args.dataset_name == "random-image": + assert not tokenize_prompt, "random-image does not support --tokenize-prompt" + input_requests = sample_random_image_requests( + num_requests=args.num_prompts, + num_images=args.random_image_num_images, + input_len=args.random_input_len, + output_len=args.random_output_len, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + apply_chat_template=args.apply_chat_template, + image_resolution=args.random_image_resolution, + ) + elif args.dataset_name == "generated-shared-prefix": + assert not tokenize_prompt + input_requests = sample_generated_shared_prefix_requests( + num_groups=args.gsp_num_groups, + prompts_per_group=args.gsp_prompts_per_group, + system_prompt_len=args.gsp_system_prompt_len, + question_len=args.gsp_question_len, + output_len=args.gsp_output_len, + tokenizer=tokenizer, + args=args, + ) + elif args.dataset_name == "mmmu": + assert not tokenize_prompt + input_requests = sample_mmmu_requests( + num_requests=args.num_prompts, + tokenizer=tokenizer, + fixed_output_len=args.random_output_len, + apply_chat_template=args.apply_chat_template, + random_sample=True, + ) + elif args.dataset_name == "mooncake": + # For mooncake, we don't generate the prompts here. + # We just load the raw trace data. The async generator will handle the rest. + if not args.dataset_path: + local_path = os.path.join("/tmp", args.mooncake_workload + "_trace.jsonl") + else: + local_path = args.dataset_path + + if not os.path.exists(local_path): + download_and_cache_file( + MOONCAKE_DATASET_URL[args.mooncake_workload], local_path + ) + + with open(local_path, "r") as f: + all_requests_data = [json.loads(line) for line in f if line.strip()] + + # Limit the number of requests based on --num-prompts + input_requests = all_requests_data[: args.num_prompts] + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + return input_requests + + +ASYNC_REQUEST_FUNCS = { + "sglang": async_request_sglang_generate, + "sglang-native": async_request_sglang_generate, + "sglang-oai": async_request_openai_completions, + "sglang-oai-chat": async_request_openai_chat_completions, + "vllm": async_request_openai_completions, + "vllm-chat": async_request_openai_chat_completions, + "lmdeploy": async_request_openai_completions, + "lmdeploy-chat": async_request_openai_chat_completions, + "trt": async_request_trt_llm, + "gserver": async_request_gserver, + "truss": async_request_truss, +} + + +@dataclass +class BenchmarkMetrics: + completed: int + total_input: int + total_output: int + total_output_retokenized: int + request_throughput: float + input_throughput: float + output_throughput: float + output_throughput_retokenized: float + total_throughput: float + total_throughput_retokenized: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + p99_ttft_ms: float + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + p99_tpot_ms: float + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + p95_itl_ms: float + p99_itl_ms: float + max_itl_ms: float + mean_e2e_latency_ms: float + median_e2e_latency_ms: float + std_e2e_latency_ms: float + p99_e2e_latency_ms: float + concurrency: float + + +SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" +MOONCAKE_DATASET_URL = { + "mooncake": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl", + "conversation": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl", + "synthetic": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/synthetic_trace.jsonl", + "toolagent": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/toolagent_trace.jsonl", +} + + +def download_and_cache_file(url: str, filename: Optional[str] = None): + """Read and cache a file from a url.""" + if filename is None: + filename = os.path.join("/tmp", url.split("/")[-1]) + + # Check if the cache file already exists + if is_file_valid_json(filename): + return filename + + print(f"Downloading from {url} to {filename}") + + # Stream the response to show the progress bar + response = requests.get(url, stream=True) + response.raise_for_status() # Check for request errors + + # Total size of the file in bytes + total_size = int(response.headers.get("content-length", 0)) + chunk_size = 1024 # Download in chunks of 1KB + + # Use tqdm to display the progress bar + with open(filename, "wb") as f, tqdm( + desc=filename, + total=total_size, + unit="B", + unit_scale=True, + unit_divisor=1024, + ) as bar: + for chunk in response.iter_content(chunk_size=chunk_size): + f.write(chunk) + bar.update(len(chunk)) + + return filename + + +def is_file_valid_json(path): + if not os.path.isfile(path): + return False + + # TODO can fuse into the real file open later + try: + with open(path) as f: + json.load(f) + return True + except JSONDecodeError as e: + print( + f"{path} exists but json loading fails ({e=}), thus treat as invalid file" + ) + return False + + +@dataclass +class DatasetRow: + prompt: str + prompt_len: int + output_len: int + image_data: Optional[List[str]] = None + timestamp: Optional[float] = None + + +async def get_mooncake_request_over_time( + input_requests: List[Dict], + tokenizer: PreTrainedTokenizerBase, + slowdown_factor: float, + num_rounds: int, +) -> AsyncGenerator[DatasetRow, None]: + """ + An async generator that yields requests based on the timestamps in the Mooncake trace file, + with support for multi-round sessions. + """ + if not input_requests: + return + + input_requests.sort(key=lambda r: r["timestamp"]) + + start_time = time.perf_counter() + trace_start_time_ms = input_requests[0]["timestamp"] + + for record in input_requests: + # Calculate when this entire session should start + relative_arrival_time_s = (record["timestamp"] - trace_start_time_ms) / 1000.0 + target_arrival_time_s = relative_arrival_time_s * slowdown_factor + + current_elapsed_time_s = time.perf_counter() - start_time + sleep_duration_s = target_arrival_time_s - current_elapsed_time_s + if sleep_duration_s > 0: + await asyncio.sleep(sleep_duration_s) + + # Once the session starts, generate all rounds for it as a burst + # This simulates a user engaging in a multi-turn conversation + + # Base user query constructed from hash_ids + user_query_base = "" + hash_ids = record.get("hash_ids", []) + for hash_id in hash_ids: + user_query_base += f"{hash_id}" + " ".join( + ["hi"] * 128 + ) # Shorter for multi-round + user_query_base += "Tell me a story based on this context." + + output_len_per_round = record.get("output_length", 256) + chat_history = [] + + for i in range(num_rounds): + # Add user query for the current round + chat_history.append( + {"role": "user", "content": f"Round {i+1}: {user_query_base}"} + ) + + # Form the full prompt from history + try: + full_prompt_text = tokenizer.apply_chat_template( + chat_history, tokenize=False, add_generation_prompt=True + ) + except Exception: + full_prompt_text = "\n".join( + [f"{msg['role']}: {msg['content']}" for msg in chat_history] + ) + + prompt_len = len(tokenizer.encode(full_prompt_text)) + + yield DatasetRow( + prompt=full_prompt_text, + prompt_len=prompt_len, + output_len=output_len_per_round, + ) + + # Add a placeholder assistant response for the next round's context + # We use a placeholder because we don't know the real response + placeholder_response = " ".join(["story"] * output_len_per_round) + chat_history.append({"role": "assistant", "content": placeholder_response}) + + +def sample_mmmu_requests( + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + fixed_output_len: Optional[int] = None, + apply_chat_template: bool = True, + random_sample: bool = True, +) -> List[DatasetRow]: + """ + Sample requests from the MMMU dataset using HuggingFace datasets. + + Args: + num_requests: Number of requests to sample. + tokenizer: Tokenizer to use for token counting. + fixed_output_len: If provided, use this fixed output length for all requests. + apply_chat_template: Whether to apply the chat template to the prompt. + random_sample: Whether to randomly sample or take the first N. + + Returns: + List of tuples (prompt, prompt_token_len, output_token_len). + """ + try: + import io + + import pybase64 + from datasets import load_dataset + except ImportError: + raise ImportError("Please install datasets: pip install datasets") + + print("Loading MMMU dataset from HuggingFace...") + + try: + print("Attempting to load MMMU Math dataset...") + mmmu_dataset = load_dataset("MMMU/MMMU", "Math", split="test") + print( + f"Successfully loaded MMMU Math dataset from HuggingFace with {len(mmmu_dataset)} examples" + ) + except Exception as e: + print(f"Failed to load MMMU Math dataset: {e}") + raise ValueError(f"Failed to load MMMU dataset: {e}") + + # Sample from the dataset + if len(mmmu_dataset) > num_requests: + if random_sample: + # Random sample + indices = random.sample(range(len(mmmu_dataset)), num_requests) + sample_dataset = mmmu_dataset.select(indices) + else: + # Take first N + sample_dataset = mmmu_dataset.select( + range(min(num_requests, len(mmmu_dataset))) + ) + else: + print(f"Dataset has less than {num_requests} examples, using all examples") + sample_dataset = mmmu_dataset + + print(f"Selected {len(sample_dataset)} examples for benchmarking") + + # Create prompts + filtered_dataset = [] + + for i, example in enumerate(sample_dataset): + try: + # Extract image_1 + image = example.get("image_1") + + if image is not None: + if hasattr(image, "save"): + # Convert RGBA images to RGB before encoding + if image.mode == "RGBA": + image = image.convert("RGB") + + # Encode image to base64 (save as PNG to support palette/alpha modes) + buffered = io.BytesIO() + image.save(buffered, format="PNG") + img_str = pybase64.b64encode(buffered.getvalue()).decode("utf-8") + image_data = f"data:image/png;base64,{img_str}" + else: + continue + + # Extract the question + question = example.get("question") + + # Construct the prompt + prompt = f"Question: {question}\n\nAnswer: " + if apply_chat_template: + try: + is_phi4_multimodal = ( + "phi-4-multimodal" in tokenizer.name_or_path.lower() + ) + if is_phi4_multimodal: + # <|endoftext10|> is the image token used in the phi-4-multimodal model. + content = prompt.replace("image 1", "<|endoftext10|>") + else: + content = [ + { + "type": "image_url", + "image_url": {"url": image_data}, + }, + {"type": "text", "text": prompt}, + ] + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": content, + } + ], + add_generation_prompt=True, + tokenize=False, + ) + except Exception as e: + # Note (Xinyuan): This is a workaround for an issue where some tokenizers do not support content as a list. (e.g. InternVL) + print( + f"Error applying chat template: {e}, fallback to tag" + ) + prompt = f"{prompt}" + + # Calculate token lengths for text only (without image data) + prompt_token_ids = tokenizer.encode(prompt) + prompt_len = len(prompt_token_ids) + + output_len = fixed_output_len if fixed_output_len is not None else 256 + + filtered_dataset.append( + DatasetRow( + prompt=prompt, + prompt_len=prompt_len, + output_len=output_len, + image_data=[image_data], + ) + ) + + except Exception as e: + print(f"Error processing example {i}: {e}") + + print(f"\nCreated {len(filtered_dataset)} MMMU prompts") + return filtered_dataset + + +def sample_sharegpt_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + fixed_output_len: Optional[int] = None, + context_len: Optional[int] = None, + prompt_suffix: Optional[str] = "", + apply_chat_template=False, +) -> List[DatasetRow]: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Download sharegpt if necessary + if not is_file_valid_json(dataset_path) and dataset_path == "": + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + + # Filter out the conversations with less than 2 turns. + dataset = [ + data + for data in dataset + if len(data.get("conversations", data.get("conversation", []))) >= 2 + ] + # Only keep the first two turns of each conversation. + dataset = [ + ( + data.get("conversations", data.get("conversation", []))[0]["value"], + data.get("conversations", data.get("conversation", []))[1]["value"], + ) + for data in dataset + ] + + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: List[DatasetRow] = [] + for i in range(len(dataset)): + if len(filtered_dataset) == num_requests: + break + + # Tokenize the prompts and completions. + prompt = dataset[i][0] + if prompt_suffix: + prompt = ( + remove_suffix(prompt, ASSISTANT_SUFFIX) + + prompt_suffix + + ASSISTANT_SUFFIX + ) + + if apply_chat_template: + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + ) + prompt = prompt.replace(tokenizer.bos_token, "") + + prompt_token_ids = tokenizer.encode(prompt) + completion = dataset[i][1] + completion_token_ids = tokenizer.encode(completion) + prompt_len = len(prompt_token_ids) + output_len = ( + len(completion_token_ids) if fixed_output_len is None else fixed_output_len + ) + + if prompt_len < 2 or output_len < 2: + # Prune too short sequences. + continue + + if context_len and prompt_len + output_len > context_len: + # Prune too long sequences. + continue + + filtered_dataset.append( + DatasetRow(prompt=prompt, prompt_len=prompt_len, output_len=output_len) + ) + + print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") + print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") + return filtered_dataset + + +def sample_random_requests( + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + dataset_path: str, + random_sample: bool = True, + return_text: bool = True, +) -> List[DatasetRow]: + input_lens = np.random.randint( + max(int(input_len * range_ratio), 1), + input_len + 1, + size=num_prompts, + ) + output_lens = np.random.randint( + int(output_len * range_ratio), + output_len + 1, + size=num_prompts, + ) + + if random_sample: + # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens + + # Download sharegpt if necessary + if not is_file_valid_json(dataset_path): + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + dataset = [ + data + for data in dataset + if len(data.get("conversations", data.get("conversation", []))) >= 2 + ] + # Only keep the first two turns of each conversation. + dataset = [ + ( + data.get("conversations", data.get("conversation", []))[0]["value"], + data.get("conversations", data.get("conversation", []))[1]["value"], + ) + for data in dataset + ] + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + input_requests: List[DatasetRow] = [] + for data in dataset: + i = len(input_requests) + if i == num_prompts: + break + + # Tokenize the prompts and completions. + prompt = data[0] + prompt_token_ids = tokenizer.encode(prompt) + prompt_len = len(prompt_token_ids) + + # Skip empty prompt + if prompt_len == 0: + continue + + if prompt_len > input_lens[i]: + input_ids = prompt_token_ids[: input_lens[i]] + else: + ratio = (input_lens[i] + prompt_len - 1) // prompt_len + input_ids = (prompt_token_ids * ratio)[: input_lens[i]] + input_content = input_ids + if return_text: + input_content = tokenizer.decode(input_content) + input_requests.append( + DatasetRow( + prompt=input_content, + prompt_len=int(input_lens[i]), + output_len=int(output_lens[i]), + ) + ) + else: + # Sample token ids from random integers. This can cause some NaN issues. + offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts) + input_requests = [] + for i in range(num_prompts): + input_content = [ + (offsets[i] + i + j) % tokenizer.vocab_size + for j in range(input_lens[i]) + ] + if return_text: + input_content = tokenizer.decode(input_content) + input_requests.append( + DatasetRow( + prompt=input_content, + prompt_len=int(input_lens[i]), + output_len=int(output_lens[i]), + ) + ) + + print(f"#Input tokens: {np.sum(input_lens)}") + print(f"#Output tokens: {np.sum(output_lens)}") + return input_requests + + +def parse_random_image_resolution(image_resolution: str) -> Tuple[int, int]: + """Parse image resolution into (width, height). + + Supports presets '1080p', '720p', '360p' and custom 'heightxwidth' format + (e.g., '1080x1920' means height=1080, width=1920). + """ + resolution_to_size = { + "4k": (3840, 2160), + "1080p": (1920, 1080), + "720p": (1280, 720), + "360p": (640, 360), + } + if image_resolution in resolution_to_size: + return resolution_to_size[image_resolution] + + res = image_resolution.strip().lower() + if "x" in res: + parts = res.split("x") + if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit(): + height = int(parts[0]) + width = int(parts[1]) + if height > 0 and width > 0: + return (width, height) + + raise ValueError( + f"Unsupported random-image resolution: {image_resolution}. " + "Choose from 4k, 1080p, 720p, 360p, or provide custom 'heightxwidth' (e.g., 1080x1920)." + ) + + +def sample_random_image_requests( + num_requests: int, + num_images: int, + input_len: int, + output_len: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + apply_chat_template: bool = True, + image_resolution: str = "1080p", +) -> List[DatasetRow]: + """Generate requests with random images. + + - Each request includes ``num_images`` random images. + - Supported resolutions: 4k (3840x2160), 1080p (1920x1080), 720p (1280x720), 360p (640x360), + or custom 'heightxwidth' (e.g., 1080x1920). + - Text lengths follow the 'random' dataset sampling rule. ``prompt_len`` + only counts text tokens and excludes image data. + """ + try: + import pybase64 + from PIL import Image + except ImportError as e: + raise ImportError( + "Please install Pillow to generate random images: pip install pillow" + ) from e + + # Parse resolution (supports presets and 'heightxwidth') + width, height = parse_random_image_resolution(image_resolution) + + # Check for potentially problematic combinations and warn user + if width * height >= 1920 * 1080 and num_images * num_requests >= 100: + warnings.warn( + f"High resolution ({width}x{height}) with {num_images * num_requests} total images " + f"may take a long time. Consider reducing resolution or image count.", + UserWarning, + stacklevel=2, + ) + + # Sample text lengths + input_lens = np.random.randint( + max(int(input_len * range_ratio), 1), input_len + 1, size=num_requests + ) + output_lens = np.random.randint( + int(output_len * range_ratio), output_len + 1, size=num_requests + ) + + def _gen_random_image_data_uri(width: int = width, height: int = height) -> str: + arr = (np.random.rand(height, width, 3) * 255).astype(np.uint8) + img = Image.fromarray(arr, mode="RGB") + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=85) + encoded = pybase64.b64encode(buf.getvalue()).decode("utf-8") + return f"data:image/jpeg;base64,{encoded}" + + dataset: List[DatasetRow] = [] + for i in range(num_requests): + # Generate text prompt + text_prompt = gen_prompt(tokenizer, int(input_lens[i])) + + # Generate image list + images = [_gen_random_image_data_uri() for _ in range(num_images)] + + prompt_str = text_prompt + if apply_chat_template: + try: + content_items = [ + {"type": "image_url", "image_url": {"url": img_url}} + for img_url in images + ] + content_items.append({"type": "text", "text": text_prompt}) + prompt_str = tokenizer.apply_chat_template( + [{"role": "user", "content": content_items}], + add_generation_prompt=True, + tokenize=False, + ) + except Exception: + # Some tokenizers do not support list content; fall back to a placeholder in the text + prompt_str = f"{text_prompt}" + + prompt_token_ids = tokenizer.encode(prompt_str) + prompt_token_len = len(prompt_token_ids) + + dataset.append( + DatasetRow( + prompt=prompt_str, + prompt_len=prompt_token_len, + output_len=int(output_lens[i]), + image_data=images, + ) + ) + + print(f"#Input tokens: {np.sum([x.prompt_len for x in dataset])}") + print(f"#Output tokens: {np.sum([x.output_len for x in dataset])}") + return dataset + + +def gen_prompt(tokenizer, token_num): + """Generate a random prompt of specified token length using tokenizer vocabulary.""" + all_available_tokens = list(tokenizer.get_vocab().values()) + selected_tokens = random.choices(all_available_tokens, k=token_num) + return tokenizer.decode(selected_tokens) + + +def get_gen_prefix_cache_path(args, tokenizer): + """Create cache directory under ~/.cache/sglang/benchmark""" + cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" + + # Create a unique cache filename based on the generation parameters + cache_key = ( + f"gen_shared_prefix_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_" + f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_output_len}_" + f"{tokenizer.__class__.__name__}.pkl" + ) + return cache_dir / cache_key + + +def sample_generated_shared_prefix_requests( + num_groups: int, + prompts_per_group: int, + system_prompt_len: int, + question_len: int, + output_len: int, + tokenizer: PreTrainedTokenizerBase, + args: argparse.Namespace, +) -> List[DatasetRow]: + """Generate benchmark requests with shared system prompts using random tokens and caching.""" + cache_path = get_gen_prefix_cache_path(args, tokenizer) + + # Try to load from cache first + if cache_path.exists(): + print(f"\nLoading cached generated input data from {cache_path}") + with open(cache_path, "rb") as f: + return pickle.load(f) + + print("\nGenerating new input data...") + + # Generate system prompts for each group + system_prompts = [] + for _ in range(num_groups): + system_prompt = gen_prompt(tokenizer, system_prompt_len) + system_prompts.append(system_prompt) + + # Generate questions + questions = [] + for _ in range(num_groups * prompts_per_group): + question = gen_prompt(tokenizer, question_len) + questions.append(question) + + # Combine system prompts with questions + input_requests = [] + total_input_tokens = 0 + total_output_tokens = 0 + + for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): + system_prompt = system_prompts[group_idx] + for prompt_idx in tqdm( + range(prompts_per_group), desc="Generating questions", leave=False + ): + question = questions[group_idx * prompts_per_group + prompt_idx] + full_prompt = f"{system_prompt}\n\n{question}" + prompt_len = len(tokenizer.encode(full_prompt)) + + input_requests.append( + DatasetRow( + prompt=full_prompt, prompt_len=prompt_len, output_len=output_len + ) + ) + total_input_tokens += prompt_len + total_output_tokens += output_len + + # Shuffle questions + random.shuffle(input_requests) + + # Print statistics + print(f"\nGenerated shared prefix dataset statistics:") + print(f"Number of groups: {num_groups}") + print(f"Prompts per group: {prompts_per_group}") + print(f"Total prompts: {len(input_requests)}") + print(f"Total input tokens: {total_input_tokens}") + print(f"Total output tokens: {total_output_tokens}") + print( + f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens" + ) + print( + f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n" + ) + + # Save to cache + cache_path.parent.mkdir(parents=True, exist_ok=True) + print(f"Caching generated input data to {cache_path}") + with open(cache_path, "wb") as f: + pickle.dump(input_requests, f) + + return input_requests + + +async def get_request( + input_requests: List[DatasetRow], + request_rate: float, + use_trace_timestamps: bool = False, + slowdown_factor: float = 1.0, +) -> AsyncGenerator[DatasetRow, None]: + if use_trace_timestamps: + print( + f"Using trace timestamps for request generation with slowdown factor {slowdown_factor}." + ) + # Sort requests by timestamp for correct replay + input_requests.sort(key=lambda r: r.timestamp) + + start_time = time.perf_counter() + trace_start_time_ms = input_requests[0].timestamp if input_requests else 0 + + for request in input_requests: + trace_time_s = (request.timestamp - trace_start_time_ms) / 1000.0 + target_arrival_time = start_time + (trace_time_s * slowdown_factor) + + sleep_duration = target_arrival_time - time.perf_counter() + if sleep_duration > 0: + await asyncio.sleep(sleep_duration) + + yield request + else: + input_requests_iter = iter(input_requests) + for request in input_requests_iter: + yield request + + if request_rate == float("inf"): + # If the request rate is infinity, then we don't need to wait. + continue + + # Sample the request interval from the exponential distribution. + interval = np.random.exponential(1.0 / request_rate) + # The next request will be sent after the interval. + await asyncio.sleep(interval) + + +def calculate_metrics( + input_requests: List[DatasetRow], + outputs: List[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase, + backend: str, +) -> Tuple[BenchmarkMetrics, List[int]]: + output_lens: List[int] = [] + retokenized_output_lens: List[int] = [] + total_input = 0 + completed = 0 + itls: List[float] = [] + tpots: List[float] = [] + ttfts: List[float] = [] + e2e_latencies: List[float] = [] + for i in range(len(outputs)): + if outputs[i].success: + output_len = outputs[i].output_len + output_lens.append(output_len) + retokenized_output_len = len( + tokenizer.encode(outputs[i].generated_text, add_special_tokens=False) + ) + retokenized_output_lens.append(retokenized_output_len) + total_input += outputs[i].prompt_len + if output_len > 1: + tpots.append((outputs[i].latency - outputs[i].ttft) / (output_len - 1)) + itls += outputs[i].itl + ttfts.append(outputs[i].ttft) + + e2e_latencies.append(outputs[i].latency) + + completed += 1 + else: + output_lens.append(0) + retokenized_output_lens.append(0) + + if completed == 0: + warnings.warn( + "All requests failed. This is likely due to a misconfiguration " + "on the benchmark arguments.", + stacklevel=2, + ) + metrics = BenchmarkMetrics( + completed=completed, + total_input=total_input, + total_output=sum(output_lens), + total_output_retokenized=sum(retokenized_output_lens), + request_throughput=completed / dur_s, + input_throughput=total_input / dur_s, + output_throughput=sum(output_lens) / dur_s, + output_throughput_retokenized=sum(retokenized_output_lens) / dur_s, + total_throughput=(total_input + sum(output_lens)) / dur_s, + total_throughput_retokenized=(total_input + sum(retokenized_output_lens)) + / dur_s, + mean_ttft_ms=np.mean(ttfts or 0) + * 1000, # ttfts is empty if streaming is not supported by backend + median_ttft_ms=np.median(ttfts or 0) * 1000, + std_ttft_ms=np.std(ttfts or 0) * 1000, + p99_ttft_ms=np.percentile(ttfts or 0, 99) * 1000, + mean_tpot_ms=np.mean(tpots or 0) * 1000, + median_tpot_ms=np.median(tpots or 0) * 1000, + std_tpot_ms=np.std(tpots or 0) * 1000, + p99_tpot_ms=np.percentile(tpots or 0, 99) * 1000, + mean_itl_ms=np.mean(itls or 0) * 1000, + median_itl_ms=np.median(itls or 0) * 1000, + std_itl_ms=np.std(itls or 0) * 1000, + p95_itl_ms=np.percentile(itls or 0, 95) * 1000, + p99_itl_ms=np.percentile(itls or 0, 99) * 1000, + max_itl_ms=np.max(itls or 0) * 1000, + mean_e2e_latency_ms=np.mean(e2e_latencies) * 1000, + median_e2e_latency_ms=np.median(e2e_latencies) * 1000, + std_e2e_latency_ms=np.std(e2e_latencies) * 1000, + p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000, + concurrency=np.sum(e2e_latencies) / dur_s, + ) + + return metrics, output_lens + + +async def benchmark( + backend: str, + api_url: str, + base_url: str, + model_id: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: List[DatasetRow], + request_rate: float, + max_concurrency: Optional[int], + disable_tqdm: bool, + lora_names: List[str], + extra_request_body: Dict[str, Any], + profile: bool, + pd_separated: bool = False, + flush_cache: bool = False, + warmup_requests: int = 1, + use_trace_timestamps: bool = False, + mooncake_slowdown_factor=1.0, + mooncake_num_rounds=1, +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + # Limit concurrency + # From https://github.com/vllm-project/vllm/pull/9390 + semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None + + async def limited_request_func(request_func_input, pbar): + if semaphore is None: + return await request_func(request_func_input=request_func_input, pbar=pbar) + async with semaphore: + return await request_func(request_func_input=request_func_input, pbar=pbar) + + # Warmup + print(f"Starting warmup with {warmup_requests} sequences...") + + # Handle the data structure difference for the warmup request + if args.dataset_name == "mooncake": + # For mooncake, input_requests is a list of dicts. + # We need to build a temporary DatasetRow for the warmup phase. + warmup_record = input_requests[0] + + # Build prompt from hash_ids, just like in the async generator + hash_ids = warmup_record.get("hash_ids", []) + prompt_text = "" + for hash_id in hash_ids: + prompt_text += f"{hash_id}" + " ".join(["hi"] * 512) + prompt_text += "Can you tell me a detailed story in 1000 words?" + + output_len = warmup_record.get("output_length", 32) + prompt_len = len(tokenizer.encode(prompt_text)) + + # Create a temporary DatasetRow object for warmup + test_request = DatasetRow( + prompt=prompt_text, + prompt_len=prompt_len, + output_len=output_len, + image_data=None, # Mooncake doesn't have image data + ) + else: + # For all other datasets, input_requests is a list of DatasetRow objects + test_request = input_requests[0] + + if lora_names is not None and len(lora_names) != 0: + lora_name = lora_names[0] + else: + lora_name = None + + # Create the test input once + test_input = RequestFuncInput( + model=model_id, + prompt=test_request.prompt, + api_url=api_url, + prompt_len=test_request.prompt_len, + output_len=min(test_request.output_len, 32), + lora_name=lora_name, + image_data=test_request.image_data, + extra_request_body=extra_request_body, + ) + + # Run warmup requests + warmup_tasks = [] + for _ in range(warmup_requests): + warmup_tasks.append( + asyncio.create_task(request_func(request_func_input=test_input)) + ) + + warmup_outputs = await asyncio.gather(*warmup_tasks) + + # Check if at least one warmup request succeeded + if warmup_requests > 0 and not any(output.success for output in warmup_outputs): + raise ValueError( + "Warmup failed - Please make sure benchmark arguments " + f"are correctly specified. Error: {warmup_outputs[0].error}" + ) + else: + print( + f"Warmup completed with {args.warmup_requests} sequences. Starting main benchmark run..." + ) + + # Flush cache + if ("sglang" in backend and _get_bool_env_var("SGLANG_IS_IN_CI")) or flush_cache: + requests.post(base_url + "/flush_cache", headers=get_auth_headers()) + + time.sleep(1.0) + + # Start profiler + if profile: + print("Starting profiler...") + profile_output = await async_request_profile( + api_url=base_url + "/start_profile" + ) + if profile_output.success: + print("Profiler started") + + # Run all requests + benchmark_start_time = time.perf_counter() + tasks: List[asyncio.Task] = [] + pbar_total = len(input_requests) + if ( + backend == "sglang" and args.dataset_name == "mooncake" + ): # Assuming mooncake is mainly for sglang or similar backends + print("Using time-based Mooncake request scheduler, ignoring --request-rate.") + request_generator = get_mooncake_request_over_time( + input_requests, tokenizer, mooncake_slowdown_factor, mooncake_num_rounds + ) + print( + f"Starting Mooncake trace replay. Sessions: {len(input_requests)}, Rounds per session: {mooncake_num_rounds}. Slowdown factor: {mooncake_slowdown_factor}" + ) + pbar_total *= args.mooncake_num_rounds + else: + request_generator = get_request(input_requests, request_rate) + + pbar = None if disable_tqdm else tqdm(total=pbar_total) + async for request in request_generator: + if lora_names is not None and len(lora_names) != 0: + idx = random.randint(0, len(lora_names) - 1) + lora_name = lora_names[idx] + else: + lora_name = None + + request_func_input = RequestFuncInput( + model=model_id, + prompt=request.prompt, + api_url=api_url, + prompt_len=request.prompt_len, + output_len=request.output_len, + lora_name=lora_name, + image_data=request.image_data, + extra_request_body=extra_request_body, + timestamp=request.timestamp, + ) + + tasks.append( + asyncio.create_task( + limited_request_func(request_func_input=request_func_input, pbar=pbar) + ) + ) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + # Stop profiler + if profile: + print("Stopping profiler...") + profile_output = await async_request_profile(api_url=base_url + "/stop_profile") + if profile_output.success: + print("Profiler stopped") + + if pbar is not None: + pbar.close() + + if "sglang" in backend: + server_info = requests.get(base_url + "/get_server_info") + if server_info.status_code == 200: + server_info_json = server_info.json() + if "decode" in server_info_json: + server_info_json = server_info_json["decode"][0] + accept_length = server_info_json["internal_states"][0].get( + "avg_spec_accept_length", None + ) + else: + accept_length = None + else: + accept_length = None + + # Compute metrics and print results + benchmark_duration = time.perf_counter() - benchmark_start_time + metrics, output_lens = calculate_metrics( + input_requests=input_requests, + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + backend=backend, + ) + + print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) + print("{:<40} {:<10}".format("Backend:", backend)) + print( + "{:<40} {:<10}".format( + "Traffic request rate:", "trace" if use_trace_timestamps else request_rate + ) + ) + print( + "{:<40} {:<10}".format( + "Max request concurrency:", + max_concurrency if max_concurrency else "not set", + ) + ) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) + print( + "{:<40} {:<10}".format( + "Total generated tokens (retokenized):", metrics.total_output_retokenized + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", metrics.request_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", metrics.input_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", metrics.output_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Total token throughput (tok/s):", metrics.total_throughput + ) + ) + print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency)) + if accept_length: + print("{:<40} {:<10.2f}".format("Accept length:", accept_length)) + print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) + print( + "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) + ) + print( + "{:<40} {:<10.2f}".format( + "Median E2E Latency (ms):", metrics.median_e2e_latency_ms + ) + ) + print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) + print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) + print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) + print("{s:{c}^{n}}".format(s="Inter-Token Latency", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) + print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) + print("{:<40} {:<10.2f}".format("P95 ITL (ms):", metrics.p95_itl_ms)) + print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) + print("{:<40} {:<10.2f}".format("Max ITL (ms):", metrics.max_itl_ms)) + print("=" * 50) + + if ( + metrics.median_ttft_ms is not None + and metrics.mean_itl_ms is not None + and metrics.output_throughput is not None + ): + result = { + # Arguments + "backend": args.backend, + "dataset_name": args.dataset_name, + "request_rate": "trace" if use_trace_timestamps else request_rate, + "max_concurrency": max_concurrency, + "sharegpt_output_len": args.sharegpt_output_len, + "random_input_len": args.random_input_len, + "random_output_len": args.random_output_len, + "random_range_ratio": args.random_range_ratio, + # Results + "duration": benchmark_duration, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + "std_e2e_latency_ms": metrics.std_e2e_latency_ms, + "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms, + "mean_ttft_ms": metrics.mean_ttft_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "std_ttft_ms": metrics.std_ttft_ms, + "p99_ttft_ms": metrics.p99_ttft_ms, + "mean_tpot_ms": metrics.mean_tpot_ms, + "median_tpot_ms": metrics.median_tpot_ms, + "std_tpot_ms": metrics.std_tpot_ms, + "p99_tpot_ms": metrics.p99_tpot_ms, + "mean_itl_ms": metrics.mean_itl_ms, + "median_itl_ms": metrics.median_itl_ms, + "std_itl_ms": metrics.std_itl_ms, + "p95_itl_ms": metrics.p95_itl_ms, + "p99_itl_ms": metrics.p99_itl_ms, + "concurrency": metrics.concurrency, + "accept_length": accept_length, + } + else: + print(f"Error running benchmark for request rate: {request_rate}") + print("-" * 30) + + # Determine output file name + if args.output_file: + output_file_name = args.output_file + else: + now = datetime.now().strftime("%m%d") + if args.dataset_name == "random-image": + output_file_name = ( + f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_" + f"{args.random_output_len}_{args.random_image_num_images}imgs_" + f"{args.random_image_resolution}.jsonl" + ) + elif args.dataset_name.startswith("random"): + output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" + else: + output_file_name = ( + f"{args.backend}_{now}_{args.num_prompts}_{args.dataset_name}.jsonl" + ) + + result_details = { + "input_lens": [output.prompt_len for output in outputs], + "output_lens": output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + } + + # Append results to a JSONL file + with open(output_file_name, "a") as file: + if args.output_details: + result_for_dump = result | result_details + else: + result_for_dump = result + file.write(json.dumps(result_for_dump) + "\n") + + return result | result_details + + +def check_chat_template(model_path): + try: + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + return "chat_template" in tokenizer.init_kwargs + except Exception as e: + print(f"Fail to load tokenizer config with error={e}") + return False + + +def set_global_args(args_: argparse.Namespace): + """Set the global args.""" + global args + args = args_ + + +def run_benchmark(args_: argparse.Namespace): + global args + args = args_ + + # Set default value for max_concurrency if not present + if not hasattr(args, "max_concurrency"): + args.max_concurrency = None + + # Set default value for warmup_requests if not present + if not hasattr(args, "warmup_requests"): + args.warmup_requests = 1 + + if not hasattr(args, "output_details"): + args.output_details = False + + if not hasattr(args, "tokenize_prompt"): + args.tokenize_prompt = False + + if not hasattr(args, "use_trace_timestamps"): + args.use_trace_timestamps = False + if not hasattr(args, "mooncake_slowdown_factor"): + args.mooncake_slowdown_factor = 1.0 + + if not hasattr(args, "mooncake_slowdown_factor"): + args.mooncake_slowdown_factor = 1.0 + + if not hasattr(args, "mooncake_num_rounds"): + args.mooncake_num_rounds = 1 + + print(f"benchmark_args={args}") + + # Set global environments + set_ulimit() + random.seed(args.seed) + np.random.seed(args.seed) + + extra_request_body = {} + if args.extra_request_body: + extra_request_body = json.loads(args.extra_request_body) + + if args.tokenize_prompt: + assert ( + args.backend == "sglang" + ), "`--tokenize-prompt` only compatible with `--backend sglang` currently" + + # Set url + if args.port is None: + args.port = { + "sglang": 30000, + "sglang-native": 30000, + "sglang-oai": 30000, + "lmdeploy": 23333, + "vllm": 8000, + "trt": 8000, + "gserver": 9988, + "truss": 8080, + }.get(args.backend, 30000) + + model_url = ( + f"{args.base_url}/v1/models" + if args.base_url + else f"http://{args.host}:{args.port}/v1/models" + ) + + if args.backend in ["sglang", "sglang-native"]: + api_url = ( + f"{args.base_url}/generate" + if args.base_url + else f"http://{args.host}:{args.port}/generate" + ) + elif args.backend in ["sglang-oai", "vllm", "lmdeploy"]: + api_url = ( + f"{args.base_url}/v1/completions" + if args.base_url + else f"http://{args.host}:{args.port}/v1/completions" + ) + elif args.backend in ["sglang-oai-chat", "vllm-chat", "lmdeploy-chat"]: + api_url = ( + f"{args.base_url}/v1/chat/completions" + if args.base_url + else f"http://{args.host}:{args.port}/v1/chat/completions" + ) + elif args.backend == "trt": + api_url = ( + f"{args.base_url}/v2/models/ensemble/generate_stream" + if args.base_url + else f"http://{args.host}:{args.port}/v2/models/ensemble/generate_stream" + ) + if args.model is None: + print("Please provide a model using `--model` when using `trt` backend.") + sys.exit(1) + elif args.backend == "gserver": + api_url = args.base_url if args.base_url else f"{args.host}:{args.port}" + args.model = args.model or "default" + elif args.backend == "truss": + api_url = ( + f"{args.base_url}/v1/models/model:predict" + if args.base_url + else f"http://{args.host}:{args.port}/v1/models/model:predict" + ) + base_url = ( + f"http://{args.host}:{args.port}" if args.base_url is None else args.base_url + ) + + # Get model name + if args.model is None: + if args.backend == "truss": + print( + "Please provide a model with `--model` when using truss backend. e.g. --model meta-llama/Llama-3.1-8B-Instruct" + ) + sys.exit(1) + try: + response = requests.get(model_url, headers=get_auth_headers()) + model_list = response.json().get("data", []) + args.model = model_list[0]["id"] if model_list else None + except Exception as e: + print(f"Failed to fetch model from {model_url}. Error: {e}") + print( + "Please specify the correct host and port using `--host` and `--port`." + ) + sys.exit(1) + + if args.model is None: + print("No model specified or found. Please provide a model using `--model`.") + sys.exit(1) + + if not check_chat_template(args.model): + print( + "\nWARNING It is recommended to use the `Chat` or `Instruct` model for benchmarking.\n" + "Because when the tokenizer counts the output tokens, if there is gibberish, it might count incorrectly.\n" + ) + + print(f"{args}\n") + + # Read dataset + backend = args.backend + model_id = args.model + tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model + tokenizer = get_tokenizer(tokenizer_id) + input_requests = get_dataset(args, tokenizer) + + # compatible with SimpleNamespace + if not hasattr(args, "flush_cache"): + args.flush_cache = False + + return asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + base_url=base_url, + model_id=model_id, + tokenizer=tokenizer, + input_requests=input_requests, + request_rate=args.request_rate, + max_concurrency=args.max_concurrency, + disable_tqdm=args.disable_tqdm, + lora_names=args.lora_name, + extra_request_body=extra_request_body, + profile=args.profile, + pd_separated=args.pd_separated, + flush_cache=args.flush_cache, + warmup_requests=args.warmup_requests, + use_trace_timestamps=args.use_trace_timestamps, + mooncake_slowdown_factor=args.mooncake_slowdown_factor, + mooncake_num_rounds=args.mooncake_num_rounds, + ) + ) + + +def set_ulimit(target_soft_limit=65535): + resource_type = resource.RLIMIT_NOFILE + current_soft, current_hard = resource.getrlimit(resource_type) + + if current_soft < target_soft_limit: + try: + resource.setrlimit(resource_type, (target_soft_limit, current_hard)) + except ValueError as e: + print(f"Fail to set RLIMIT_NOFILE: {e}") + + +class LoRAPathAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, []) + for lora_name in values: + getattr(namespace, self.dest).append(lora_name) + + +if __name__ == "__main__": + parser = ArgumentParser(description="Benchmark the online serving throughput.") + parser.add_argument( + "--backend", + type=str, + choices=list(ASYNC_REQUEST_FUNCS.keys()), + default="sglang", + help="Must specify a backend, depending on the LLM Inference Engine.", + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + parser.add_argument( + "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." + ) + parser.add_argument( + "--port", + type=int, + help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=[ + "sharegpt", + "random", + "random-ids", + "generated-shared-prefix", + "mmmu", + "random-image", + "mooncake", + ], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument( + "--dataset-path", type=str, default="", help="Path to the dataset." + ) + parser.add_argument( + "--model", + type=str, + help="Name or path of the model. If not set, the default model will request /v1/models for conf.", + ) + parser.add_argument( + "--tokenizer", + type=str, + help="Name or path of the tokenizer. If not set, using the model conf.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1000, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--sharegpt-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output length from the ShareGPT dataset.", + ) + parser.add_argument( + "--sharegpt-context-len", + type=int, + default=None, + help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=1024, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + default=1024, + type=int, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=0.0, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + # random-image dataset args + parser.add_argument( + "--random-image-num-images", + type=int, + default=1, + help="Number of images per request (only available with the random-image dataset)", + ) + parser.add_argument( + "--random-image-resolution", + type=str, + default="1080p", + help=( + "Resolution of random images for random-image dataset. " + "Supports presets 4k/1080p/720p/360p or custom 'heightxwidth' (e.g., 1080x1920)." + ), + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " + "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", + ) + parser.add_argument( + "--use-trace-timestamps", + action="store_true", + help="Use timestamps from the trace file for request scheduling. Only valid for 'mooncake' dataset.", + ) + parser.add_argument( + "--max-concurrency", + type=int, + default=None, + help="Maximum number of concurrent requests. This can be used " + "to help simulate an environment where a higher level component " + "is enforcing a maximum number of concurrent requests. While the " + "--request-rate argument controls the rate at which requests are " + "initiated, this argument will control how many are actually allowed " + "to execute at a time. This means that when used in combination, the " + "actual request rate may be lower than specified with --request-rate, " + "if the server is not processing requests fast enough to keep up.", + ) + parser.add_argument("--output-file", type=str, help="Output JSONL file name.") + parser.add_argument( + "--output-details", action="store_true", help="Output details of benchmarking." + ) + parser.add_argument( + "--disable-tqdm", + action="store_true", + help="Specify to disable tqdm progress bar.", + ) + parser.add_argument( + "--disable-stream", + action="store_true", + help="Disable streaming mode.", + ) + parser.add_argument( + "--return-logprob", + action="store_true", + help="Return logprob.", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--disable-ignore-eos", + action="store_true", + help="Disable ignoring EOS.", + ) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + parser.add_argument( + "--apply-chat-template", + action="store_true", + help="Apply chat template", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "SGLANG_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--lora-name", + type=str, + nargs="*", + default=None, + action=LoRAPathAction, + help="The names of LoRA adapters. You can provide a list of names in the format {name} {name} {name}...", + ) + parser.add_argument( + "--prompt-suffix", + type=str, + default="", + help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.", + ) + parser.add_argument( + "--pd-separated", + action="store_true", + help="Benchmark PD disaggregation server", + ) + parser.add_argument( + "--flush-cache", + action="store_true", + help="Flush the cache before running the benchmark", + ) + parser.add_argument( + "--warmup-requests", + type=int, + default=1, + help="Number of warmup requests to run before the benchmark", + ) + parser.add_argument( + "--tokenize-prompt", + action="store_true", + help="Use integer ids instead of string for inputs. Useful to control prompt lengths accurately", + ) + + group = parser.add_argument_group("generated-shared-prefix dataset arguments") + group.add_argument( + "--gsp-num-groups", + type=int, + default=64, + help="Number of system prompt groups for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-prompts-per-group", + type=int, + default=16, + help="Number of prompts per system prompt group for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-system-prompt-len", + type=int, + default=2048, + help="Target length in tokens for system prompts in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-question-len", + type=int, + default=128, + help="Target length in tokens for questions in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-output-len", + type=int, + default=256, + help="Target length in tokens for outputs in generated-shared-prefix dataset", + ) + mooncake_group = parser.add_argument_group("mooncake dataset arguments") + mooncake_group.add_argument( + "--mooncake-slowdown-factor", + type=float, + default=1.0, + help="Slowdown factor for replaying the mooncake trace. " + "A value of 2.0 means the replay is twice as slow. " + "NOTE: --request-rate is IGNORED in mooncake mode.", + ) + mooncake_group.add_argument( + "--mooncake-num-rounds", + type=int, + default=1, + help="Number of conversation rounds for each session in the mooncake dataset. " + "A value > 1 will enable true multi-turn session benchmarking.", + ) + mooncake_group.add_argument( + "--mooncake-workload", + type=str, + default="conversation", + choices=[ + "mooncake", + "conversation", + "synthetic", + "toolagent", + ], + help="Underlying workload for the mooncake dataset.", + ) + args = parser.parse_args() + run_benchmark(args) diff --git a/python/sglang/check_env.py b/python/sglang/check_env.py new file mode 100644 index 000000000..1870e3207 --- /dev/null +++ b/python/sglang/check_env.py @@ -0,0 +1,305 @@ +"""Check environment configurations and dependency versions.""" + +import importlib.metadata +import os +import resource +import subprocess +import sys +from collections import OrderedDict, defaultdict + +import torch + +from sglang.srt.utils import is_hip + + +def is_cuda_v2(): + return torch.version.cuda is not None + + +# List of packages to check versions +PACKAGE_LIST = [ + "sglang", + "sgl_kernel", + "flashinfer_python", + "triton", + "transformers", + "torchao", + "numpy", + "aiohttp", + "fastapi", + "hf_transfer", + "huggingface_hub", + "interegular", + "modelscope", + "orjson", + "outlines", + "packaging", + "psutil", + "pydantic", + "python-multipart", + "pyzmq", + "torchao", + "uvicorn", + "uvloop", + "vllm", + "xgrammar", + "openai", + "tiktoken", + "anthropic", + "litellm", + "decord", +] + + +def get_package_versions(packages): + """ + Get versions of specified packages. + """ + versions = {} + for package in packages: + package_name = package.split("==")[0].split(">=")[0].split("<=")[0] + try: + version = importlib.metadata.version(package_name) + versions[package_name] = version + except ModuleNotFoundError: + versions[package_name] = "Module Not Found" + return versions + + +def get_cuda_info(): + """ + Get CUDA-related information if available. + """ + if is_cuda_v2(): + cuda_info = {"CUDA available": torch.cuda.is_available()} + + if cuda_info["CUDA available"]: + cuda_info.update(_get_gpu_info()) + cuda_info.update(_get_cuda_version_info()) + + return cuda_info + elif is_hip(): + cuda_info = {"ROCM available": torch.cuda.is_available()} + + if cuda_info["ROCM available"]: + cuda_info.update(_get_gpu_info()) + cuda_info.update(_get_cuda_version_info()) + + return cuda_info + + +def _get_gpu_info(): + """ + Get information about available GPUs. + """ + devices = defaultdict(list) + capabilities = defaultdict(list) + for k in range(torch.cuda.device_count()): + devices[torch.cuda.get_device_name(k)].append(str(k)) + capability = torch.cuda.get_device_capability(k) + capabilities[f"{capability[0]}.{capability[1]}"].append(str(k)) + + gpu_info = {} + for name, device_ids in devices.items(): + gpu_info[f"GPU {','.join(device_ids)}"] = name + + if len(capabilities) == 1: + # All GPUs have the same compute capability + cap, gpu_ids = list(capabilities.items())[0] + gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap + else: + # GPUs have different compute capabilities + for cap, gpu_ids in capabilities.items(): + gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap + + return gpu_info + + +def _get_cuda_version_info(): + """ + Get CUDA version information. + """ + if is_cuda_v2(): + from torch.utils.cpp_extension import CUDA_HOME + + cuda_info = {"CUDA_HOME": CUDA_HOME} + + if CUDA_HOME and os.path.isdir(CUDA_HOME): + cuda_info.update(_get_nvcc_info()) + cuda_info.update(_get_cuda_driver_version()) + + return cuda_info + elif is_hip(): + from torch.utils.cpp_extension import ROCM_HOME as ROCM_HOME + + cuda_info = {"ROCM_HOME": ROCM_HOME} + + if ROCM_HOME and os.path.isdir(ROCM_HOME): + cuda_info.update(_get_nvcc_info()) + cuda_info.update(_get_cuda_driver_version()) + + return cuda_info + else: + cuda_info = {"CUDA_HOME": ""} + return cuda_info + + +def _get_nvcc_info(): + """ + Get NVCC version information. + """ + if is_cuda_v2(): + from torch.utils.cpp_extension import CUDA_HOME + + try: + nvcc = os.path.join(CUDA_HOME, "bin/nvcc") + nvcc_output = ( + subprocess.check_output(f'"{nvcc}" -V', shell=True) + .decode("utf-8") + .strip() + ) + return { + "NVCC": nvcc_output[ + nvcc_output.rfind("Cuda compilation tools") : nvcc_output.rfind( + "Build" + ) + ].strip() + } + except subprocess.SubprocessError: + return {"NVCC": "Not Available"} + elif is_hip(): + from torch.utils.cpp_extension import ROCM_HOME + + try: + hipcc = os.path.join(ROCM_HOME, "bin/hipcc") + hipcc_output = ( + subprocess.check_output(f'"{hipcc}" --version', shell=True) + .decode("utf-8") + .strip() + ) + return { + "HIPCC": hipcc_output[ + hipcc_output.rfind("HIP version") : hipcc_output.rfind("AMD clang") + ].strip() + } + except subprocess.SubprocessError: + return {"HIPCC": "Not Available"} + else: + return {"NVCC": "Not Available"} + + +def _get_cuda_driver_version(): + """ + Get CUDA driver version. + """ + versions = set() + if is_cuda_v2(): + try: + output = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader,nounits", + ] + ) + versions = set(output.decode().strip().split("\n")) + if len(versions) == 1: + return {"CUDA Driver Version": versions.pop()} + else: + return {"CUDA Driver Versions": ", ".join(sorted(versions))} + except subprocess.SubprocessError: + return {"CUDA Driver Version": "Not Available"} + elif is_hip(): + try: + output = subprocess.check_output( + [ + "rocm-smi", + "--showdriverversion", + "--csv", + ] + ) + versions = set(output.decode().strip().split("\n")) + versions.discard("name, value") + ver = versions.pop() + ver = ver.replace('"Driver version", ', "").replace('"', "") + + return {"ROCM Driver Version": ver} + except subprocess.SubprocessError: + return {"ROCM Driver Version": "Not Available"} + else: + return {"CUDA Driver Version": "Not Available"} + + +def get_gpu_topology(): + """ + Get GPU topology information. + """ + if is_cuda_v2(): + try: + result = subprocess.run( + ["nvidia-smi", "topo", "-m"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return "\n" + result.stdout if result.returncode == 0 else None + except subprocess.SubprocessError: + return None + elif is_hip(): + try: + result = subprocess.run( + ["rocm-smi", "--showtopotype"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return "\n" + result.stdout if result.returncode == 0 else None + except subprocess.SubprocessError: + return None + else: + return None + + +def get_hypervisor_vendor(): + try: + output = subprocess.check_output(["lscpu"], text=True) + for line in output.split("\n"): + if "Hypervisor vendor:" in line: + return line.split(":")[1].strip() + return None + except: + return None + + +def check_env(): + """ + Check and print environment information. + """ + env_info = OrderedDict() + env_info["Python"] = sys.version.replace("\n", "") + env_info.update(get_cuda_info()) + env_info["PyTorch"] = torch.__version__ + env_info.update(get_package_versions(PACKAGE_LIST)) + + gpu_topo = get_gpu_topology() + if gpu_topo: + if is_cuda_v2(): + env_info["NVIDIA Topology"] = gpu_topo + elif is_hip(): + env_info["AMD Topology"] = gpu_topo + + hypervisor_vendor = get_hypervisor_vendor() + if hypervisor_vendor: + env_info["Hypervisor vendor"] = hypervisor_vendor + + ulimit_soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE) + env_info["ulimit soft"] = ulimit_soft + + for k, v in env_info.items(): + print(f"{k}: {v}") + + +if __name__ == "__main__": + check_env() diff --git a/python/sglang/compile_deep_gemm.py b/python/sglang/compile_deep_gemm.py new file mode 100644 index 000000000..e59036f7b --- /dev/null +++ b/python/sglang/compile_deep_gemm.py @@ -0,0 +1,184 @@ +""" +Compile DeepGEMM Kernels for a model with specify server arguments + +This script launches a server for capturing DeepGEMM calls and then compiles the kernels. +It accepts server arguments (the same as launch_server.py). + +Usage: +python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code + +""" + +import argparse +import dataclasses +import multiprocessing +import os +import time + +import requests + +from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.managers.tokenizer_manager import TokenizerManager +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import kill_process_tree +from sglang.srt.warmup import warmup + +multiprocessing.set_start_method("spawn", force=True) + +# Reduce warning +os.environ["SGL_IN_DEEPGEMM_PRECOMPILE_STAGE"] = "1" +# Force enable deep gemm +os.environ["SGL_ENABLE_JIT_DEEPGEMM"] = "1" +# Force enable mha chunked kv for DeepSeek V3 to avoid missing kv_b_proj DeepGEMM case +os.environ["SGL_CHUNKED_PREFIX_CACHE_THRESHOLD"] = "0" + + +@dataclasses.dataclass +class CompileArgs: + timeout: int = 3600 + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--timeout", type=int, default=CompileArgs.timeout) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + # use the default value's type to cast the args into correct types. + attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] + return cls( + **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} + ) + + +@warmup("compile-deep-gemm") +async def warm_up_compile( + disaggregation_mode: str, tokenizer_manager: TokenizerManager +): + print("\nGenerate warm up request for compiling DeepGEMM...\n") + generate_req_input = GenerateReqInput( + input_ids=[0, 1, 2, 3], + sampling_params={ + "temperature": 0.0, + "max_new_tokens": 8, + "ignore_eos": True, + }, + ) + if disaggregation_mode != "null": + generate_req_input.bootstrap_room = 0 + generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST + + await tokenizer_manager.generate_request(generate_req_input, None).__anext__() + + +def launch_server_internal(server_args): + try: + launch_server(server_args) + except Exception as e: + raise e + finally: + kill_process_tree(os.getpid(), include_parent=False) + + +def launch_server_process_and_send_one_request( + server_args: ServerArgs, compile_args: CompileArgs +): + proc = multiprocessing.Process(target=launch_server_internal, args=(server_args,)) + proc.start() + base_url = f"http://{server_args.host}:{server_args.port}" + timeout = compile_args.timeout + + start_time = time.perf_counter() + while time.perf_counter() - start_time < timeout: + try: + headers = { + "Content-Type": "application/json; charset=utf-8", + } + if server_args.node_rank == 0: + response = requests.get(f"{base_url}/v1/models", headers=headers) + else: + # This http api is created by launch_dummy_health_check_server for none-rank0 node. + response = requests.get(f"{base_url}/health", headers=headers) + if response.status_code == 200: + # Rank-0 node send a request to sync with other node and then return. + if server_args.node_rank == 0: + response = requests.post( + f"{base_url}/generate", + json={ + "input_ids": [0, 1, 2, 3], + "sampling_params": { + "max_new_tokens": 8, + "temperature": 0, + }, + }, + timeout=600, + ) + if response.status_code != 200: + error = response.json() + raise RuntimeError(f"Sync request failed: {error}") + # Other nodes should wait for the exit signal from Rank-0 node. + else: + start_time_waiting = time.perf_counter() + while proc.is_alive(): + if time.perf_counter() - start_time_waiting < timeout: + time.sleep(10) + else: + raise TimeoutError("Waiting for main node timeout!") + return proc + except requests.RequestException: + pass + time.sleep(10) + raise TimeoutError( + "DeepGEMM Kernels compilation timeout." + "\n\nFeel free and please restart the command." + ) + + +def refine_server_args(server_args: ServerArgs, compile_args: CompileArgs): + # Disable cuda graph and torch compile to save time + server_args.disable_cuda_graph = True + server_args.enable_torch_compile = False + print(f"Disable CUDA Graph and Torch Compile to save time...") + + # Set watchdog timeout to compile_args.timeout because compilation will take a long time + server_args.watchdog_timeout = compile_args.timeout + server_args.warmups = "compile-deep-gemm" + + +def run_compile(server_args: ServerArgs, compile_args: CompileArgs): + print( + "Begin DeepGEMM Kernels compilation...\n" + "It may take a long time and timeout maybe raised " + "while the compilation is still in progress.\n" + "Just feel free to restart the command " + "until the compilation is fully finished.\n" + ) + + proc = launch_server_process_and_send_one_request(server_args, compile_args) + + print("\nDeepGEMM Kernels compilation finished successfully.") + + # Sleep for safety + time.sleep(10) + if proc.is_alive(): + # This is the rank0 node. + kill_process_tree(proc.pid) + else: + try: + kill_process_tree(proc.pid) + except Exception: + pass + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + CompileArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + compile_args = CompileArgs.from_cli_args(args) + + refine_server_args(server_args, compile_args) + + run_compile(server_args, compile_args) diff --git a/python/sglang/eval/llama3_eval.py b/python/sglang/eval/llama3_eval.py new file mode 100644 index 000000000..253cdf275 --- /dev/null +++ b/python/sglang/eval/llama3_eval.py @@ -0,0 +1,315 @@ +# Adapt from https://github.com/fw-ai/llm_eval_meta + +import argparse +import asyncio +import os +import pickle +import re +import shutil +from collections import defaultdict +from dataclasses import dataclass + +import httpx +import numpy as np +import openai +from datasets import load_dataset +from openai import AsyncOpenAI +from tqdm import tqdm + +# Mapping providers to their clients and models +provider_to_models = { + "b10": { + "8b": "meta-llama/Llama-3.1-8B-Instruct", + "70b": "meta-llama/Llama-3.1-70B-Instruct", + "405b": "meta-llama/Llama-3.1-405B-Instruct", + }, + "oai": { + "8b": "meta-llama/Llama-3.1-8B-Instruct", + "70b": "meta-llama/Llama-3.1-70B-Instruct", + "405b": "meta-llama/Llama-3.1-405B-Instruct", + }, + "sgl": { + "8b": "meta-llama/Llama-3.1-8B-Instruct", + "70b": "meta-llama/Llama-3.1-70B-Instruct", + "405b": "meta-llama/Llama-3.1-405B-Instruct", + }, +} + + +async def fetch_responses( + client, prompt, semaphore, index, provider, model_size, output_dir, max_tokens +): + output_file = os.path.join(output_dir, f"response_{index}.pkl") + if os.path.exists(output_file): + print(f"File {output_file} already exists, skipping.") + return + + async with semaphore: + response = await client.completions.create( + model=provider_to_models[provider][model_size], + prompt=prompt, + temperature=0.0, + max_tokens=max_tokens, + ) + if isinstance(response, openai.BadRequestError): + with open(output_file, "wb") as f: + pickle.dump("bad_response", f) + assert isinstance(response, openai.types.completion.Completion) + # Save response to a file + with open(output_file, "wb") as f: + pickle.dump(response, f) + + +TASK_TO_MAX_TOKENS = { + "evals__mmlu__details": 1, + "evals__mmlu__0_shot__cot__details": 1024, + # Official meta uses 1024, but a small % (.05) of questions are answered correctly after relaxing + "evals__mmlu_pro__details": 2048, + "evals__gsm8k__details": 1024, +} + +TASK_TO_EVAL_SET = { + "mmlu": "evals__mmlu__details", + "mmlu_cot": "evals__mmlu__0_shot__cot__details", + "mmlu_pro": "evals__mmlu_pro__details", + "gsm8k": "evals__gsm8k__details", +} + + +class CustomAsyncHTTPXClient(httpx.AsyncClient): + async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response: + request.url = httpx.URL( + f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict" + ) + return await super().send(request, *args, **kwargs) + + +def get_client(provider): + if provider not in "b10": + if os.getenv("OPENAI_API_KEY") == None: + os.environ["OPENAI_API_KEY"] = "EMPTY" + return { + "oai": AsyncOpenAI(base_url="http://127.0.0.1:8000/v1/"), + "b10": AsyncOpenAI( + api_key=f"Api-Key {os.getenv('OPENAI_API_KEY')}", + base_url=f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict", + http_client=CustomAsyncHTTPXClient(), + ), + "sgl": AsyncOpenAI(base_url="http://127.0.0.1:30000/v1/"), + }[provider] + + +# Define the benchmark function +async def benchmark(args): + ds = load_dataset( + "meta-llama/Llama-3.1-405B-Instruct-evals", + f"Llama-3.1-405B-Instruct-{TASK_TO_EVAL_SET[args.task]}", + ) + semaphore = asyncio.Semaphore(args.concurrency) # Limit to 16 concurrent tasks + + if args.num_examples is None: + args.num_examples = len(ds["latest"]["input_final_prompts"]) + prompts = ds["latest"]["input_final_prompts"][: args.num_examples] + + # Create the output directory if it does not exist + os.makedirs(args.output_dir, exist_ok=True) + + tasks = [] + # Create the tasks with tqdm progress bar + max_tokens = TASK_TO_MAX_TOKENS[TASK_TO_EVAL_SET[args.task]] + client = get_client(args.provider) + for idx, prompt in enumerate(tqdm(prompts, desc="Creating tasks")): + tasks.append( + asyncio.create_task( + fetch_responses( + client, + f"<|begin_of_text|>{prompt[0]}", + semaphore, + idx, + args.provider, + args.model_size, + args.output_dir, + max_tokens=max_tokens, + ) + ) + ) + + # Run the tasks with tqdm progress bar + for future in tqdm( + asyncio.as_completed(tasks), total=len(tasks), desc="Processing tasks" + ): + await future + + +def get_mmlu_answer(response): + if response is not None: + return response.choices[0].text.lstrip().rstrip().upper().replace(".", "") + return None + + +def get_mmlu_cot_answer(response): + pattern = r"The best answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + return match.group(1).replace(".", "").replace("*", "") + + pattern = r"the best answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + return match.group(1).replace(".", "") + + pattern = r"The correct answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + return match.group(1).replace(".", "") + + pattern = r"the correct answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + return match.group(1).replace(".", "") + + +def get_answer_gsm8k(response): + pattern = r"The final answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + s = match.group(1) + for ok_symbol in ["%", "$"]: + s = s.replace(ok_symbol, "") + return s + + +TASK_TO_ANSWER_EXTRACTOR = { + "evals__mmlu__details": get_mmlu_answer, + "evals__mmlu__0_shot__cot__details": get_mmlu_cot_answer, + "evals__gsm8k__details": get_answer_gsm8k, + "evals__mmlu_pro__details": get_mmlu_cot_answer, +} + + +def get_dataset_from_task(task, response_path, model_size): + ds_405b = load_dataset( + f"meta-llama/Llama-3.1-405B-Instruct-evals", + f"Llama-3.1-405B-Instruct-{task}", + ) + ds_405b_hash_order = [x[0] for x in ds_405b["latest"]["input_final_prompts_hash"]] + + if "70b" in model_size or "8b" in model_size: + if "70" in model_size: + ref_model_ds = load_dataset( + f"meta-llama/Llama-3.1-70B-Instruct-evals", + f"Llama-3.1-70B-Instruct-{task}", + ) + else: + ref_model_ds = load_dataset( + f"meta-llama/Llama-3.1-8B-Instruct-evals", + f"Llama-3.1-8B-Instruct-{task}", + ) + + hash_to_row = {} + for row in ref_model_ds["latest"]: + hash_to_row[row["input_final_prompts_hash"][0]] = row + reordered_rows = [] + for prompt_hash in ds_405b_hash_order: + reordered_rows.append(hash_to_row[prompt_hash]) + ref_model_ds["latest"] = reordered_rows + return ref_model_ds + + return ds_405b + + +def analyze(task, response_path, model_size): + ds = get_dataset_from_task(task, response_path, model_size) + + responses = [] + total = len(ds["latest"]) + + for i in range(0, total): + response = pickle.load( + open(os.path.join(response_path, f"response_{i}.pkl"), "rb") + ) + responses.append(response) + + @dataclass + class Stats: + correct: int = 0 + total: int = 0 + meta_correct: int = 0 + + average: float = None + + subtask_name_to_stats = defaultdict(lambda: Stats()) + + for response, ds_row in zip(responses, ds["latest"]): + model_answer = TASK_TO_ANSWER_EXTRACTOR[task](response) + + subtask = ds_row["subtask_name"] + + is_eval_correct = model_answer in ds_row["input_correct_responses"] + if is_eval_correct: + subtask_name_to_stats[subtask].correct += 1 + + if ds_row["is_correct"]: + subtask_name_to_stats[subtask].meta_correct += 1 + + subtask_name_to_stats[subtask].total += 1 + + micro_stats = Stats() + for subtask, stats in subtask_name_to_stats.items(): + stats.average = stats.correct / stats.total + stats.meta_average = stats.meta_correct / stats.total + + micro_stats.correct += stats.correct + micro_stats.total += stats.total + micro_stats.meta_correct += stats.meta_correct + + micro_stats.average = micro_stats.correct / micro_stats.total + micro_stats.meta_average = micro_stats.meta_correct / micro_stats.total + + print("Macro average", np.mean([x.average for x in subtask_name_to_stats.values()])) + print( + "Meta Macro average", + np.mean([x.meta_average for x in subtask_name_to_stats.values()]), + ) + print("Micro average", micro_stats.average) + print("Meta Micro average", micro_stats.meta_average) + + +# Entry point for the script +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Script to run model with specified parameters." + ) + parser.add_argument( + "--model-size", + type=str, + default="8b", + help="Size of the model (e.g., 8b or 70b)", + ) + parser.add_argument( + "--provider", + type=str, + default="sgl", + help="Provider name (e.g., sgl, oai, b10)", + ) + parser.add_argument( + "--task", + type=str, + required=True, + help="Task (e.g., mmlu, mmlu_cot, mmlu_pro, gsm8k)", + ) + parser.add_argument( + "--num-examples", type=int, default=None, help="Number of examples to process" + ) + parser.add_argument("--concurrency", type=int, default=16) + parser.add_argument( + "--output-dir", + type=str, + default="tmp-output-dir", + help="Directory to save responses", + ) + + args = parser.parse_args() + asyncio.run(benchmark(args)) + analyze(TASK_TO_EVAL_SET[args.task], args.output_dir, args.model_size) + shutil.rmtree("tmp-output-dir", ignore_errors=True) diff --git a/python/sglang/eval/loogle_eval.py b/python/sglang/eval/loogle_eval.py new file mode 100644 index 000000000..895362cd1 --- /dev/null +++ b/python/sglang/eval/loogle_eval.py @@ -0,0 +1,164 @@ +import argparse +import asyncio +import os +import pickle +from pathlib import Path +from typing import List + +import openai +import torch +from bert_score import BERTScorer +from datasets import load_dataset +from tqdm import tqdm + + +def get_client(api_url: str) -> openai.AsyncOpenAI: + if os.getenv("OPENAI_API_KEY") is None: + os.environ["OPENAI_API_KEY"] = "EMPTY" + return openai.AsyncOpenAI(base_url=api_url) + + +def get_dataset(): + return load_dataset("bigai-nlco/LooGLE", "longdep_qa", split="test") + + +async def fetch_response( + client: openai.AsyncOpenAI, + context: str, + question: str, + semaphore: asyncio.Semaphore, + index: int, + model: str, + output_dir: Path, +): + output_file = output_dir / f"response_{index}.pkl" + if output_file.exists(): + return + + prompt = ( + "Please answer the question based on the long texts below.\n" + f"{context}\n" + f"Question: {question}\n" + "Answer:" + ) + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": prompt}, + ] + + async with semaphore: + try: + response = await client.chat.completions.create( + model=model, + messages=messages, + temperature=0.0, + max_tokens=512, + ) + except openai.BadRequestError as e: + with open(output_file, "wb") as f: + pickle.dump({"error": str(e)}, f) + return + + with open(output_file, "wb") as f: + pickle.dump(response, f) + + +async def benchmark(args): + dataset = get_dataset() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + client = get_client(args.api_url) + semaphore = asyncio.Semaphore(args.max_concurrency) + + tasks: List[asyncio.Task] = [] + for idx, ex in enumerate(dataset): + if idx >= args.num_prompts: + break + tasks.append( + asyncio.create_task( + fetch_response( + client, + ex["context"], + ex["question"], + semaphore, + idx, + args.model, + output_dir, + ) + ) + ) + + for _ in tqdm( + asyncio.as_completed(tasks), total=len(tasks), desc="Running benchmark" + ): + await _ + + +def analyse(args): + dataset = get_dataset() + output_dir = Path(args.output_dir) + + device = "cuda" if torch.cuda.is_available() else "cpu" + scorer = BERTScorer(lang="en", device=device) + + hyps: List[str] = [] + refs: List[str] = [] + for idx, ex in enumerate(tqdm(dataset, desc="Loading responses")): + if idx >= args.num_prompts: + break + pkl_file = output_dir / f"response_{idx}.pkl" + if not pkl_file.exists(): + raise FileNotFoundError(pkl_file) + + response = pickle.load(open(pkl_file, "rb")) + if isinstance(response, dict) and "error" in response: + continue + + hyps.append(response.choices[0].message.content.strip()) + refs.append(ex["answer"]) + + if not hyps: + print("No valid responses to score!") + return + + batch_size = 64 + all_f1: List[float] = [] + for i in tqdm(range(0, len(hyps), batch_size), desc="Scoring batches"): + h_batch = hyps[i : i + batch_size] + r_batch = refs[i : i + batch_size] + _, _, f1_scores = scorer.score(h_batch, r_batch, verbose=False) + all_f1.extend([float(x) for x in f1_scores]) + + avg = sum(all_f1) / len(all_f1) + print(f"Average BERTScore (F1): {avg:.2%}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Run benchmark and evaluation in one go." + ) + parser.add_argument( + "--api-url", + default="http://127.0.0.1:30000/v1", + help="OpenAI‑compatible API base URL", + ) + parser.add_argument( + "--model", + default="meta-llama/Llama-4-Maverick-17B-128E-Instruct", + help="Model name or ID, only used for model name", + ) + parser.add_argument( + "--max-concurrency", type=int, default=144, help="Maximum concurrent requests" + ) + parser.add_argument( + "--output-dir", default="tmp-output-dir", help="Directory for cached responses" + ) + parser.add_argument( + "--num-prompts", type=int, default=10000, help="Number of prompts to run" + ) + args = parser.parse_args() + + asyncio.run(benchmark(args)) + + analyse(args) diff --git a/python/sglang/global_config.py b/python/sglang/global_config.py new file mode 100644 index 000000000..f006bd94c --- /dev/null +++ b/python/sglang/global_config.py @@ -0,0 +1,53 @@ +"""Global configurations""" + +import os + + +class GlobalConfig: + """ + Store some global constants. + + See also python/sglang/srt/managers/schedule_batch.py::global_server_args_dict, which stores + many global runtime arguments as well. + """ + + def __init__(self): + # Verbosity level + # 0: do not output anything + # 2: output final text after every run + self.verbosity = 0 + + # Default backend of the language + self.default_backend = None + + # Runtime constants: New generation token ratio estimation + self.default_init_new_token_ratio = float( + os.environ.get("SGLANG_INIT_NEW_TOKEN_RATIO", 0.7) + ) + self.default_min_new_token_ratio_factor = float( + os.environ.get("SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR", 0.14) + ) + self.default_new_token_ratio_decay_steps = float( + os.environ.get("SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS", 600) + ) + self.torch_empty_cache_interval = float( + os.environ.get( + "SGLANG_EMPTY_CACHE_INTERVAL", -1 + ) # in seconds. Set if you observe high memory accumulation over a long serving period. + ) + # Runtime constants: others + self.retract_decode_steps = 20 + self.flashinfer_workspace_size = os.environ.get( + "FLASHINFER_WORKSPACE_SIZE", 384 * 1024 * 1024 + ) + + # Output tokenization configs + self.skip_special_tokens_in_output = True + self.spaces_between_special_tokens_in_out = True + + # Language frontend interpreter optimization configs + self.enable_precache_with_tracing = True + self.enable_parallel_encoding = True + + +global_config = GlobalConfig() diff --git a/python/sglang/lang/api.py b/python/sglang/lang/api.py new file mode 100644 index 000000000..a8d2e43e6 --- /dev/null +++ b/python/sglang/lang/api.py @@ -0,0 +1,286 @@ +"""Public APIs of the language.""" + +import re +from typing import Callable, List, Optional, Union + +from sglang.global_config import global_config +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.choices import ChoicesSamplingMethod, token_length_normalized +from sglang.lang.ir import ( + SglExpr, + SglExprList, + SglFunction, + SglGen, + SglImage, + SglRoleBegin, + SglRoleEnd, + SglSelect, + SglSeparateReasoning, + SglVideo, +) + + +def function( + func: Optional[Callable] = None, num_api_spec_tokens: Optional[int] = None +): + if func: + return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens) + + def decorator(func): + return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens) + + return decorator + + +def Runtime(*args, **kwargs): + # Avoid importing unnecessary dependency + from sglang.lang.backend.runtime_endpoint import Runtime + + return Runtime(*args, **kwargs) + + +def Engine(*args, **kwargs): + # Avoid importing unnecessary dependency + from sglang.srt.entrypoints.engine import Engine + + return Engine(*args, **kwargs) + + +def set_default_backend(backend: BaseBackend): + global_config.default_backend = backend + + +def flush_cache(backend: Optional[BaseBackend] = None): + backend = backend or global_config.default_backend + if backend is None: + return False + + # If backend is Runtime + if hasattr(backend, "endpoint"): + backend = backend.endpoint + return backend.flush_cache() + + +def get_server_info(backend: Optional[BaseBackend] = None): + backend = backend or global_config.default_backend + if backend is None: + return None + + # If backend is Runtime + if hasattr(backend, "endpoint"): + backend = backend.endpoint + return backend.get_server_info() + + +def gen( + name: Optional[str] = None, + max_tokens: Optional[int] = None, + min_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, + dtype: Optional[Union[type, str]] = None, + choices: Optional[List[str]] = None, + choices_method: Optional[ChoicesSamplingMethod] = None, + regex: Optional[str] = None, + json_schema: Optional[str] = None, +): + """Call the model to generate. See the meaning of the arguments in docs/backend/sampling_params.md""" + + if choices: + return SglSelect( + name, + choices, + 0.0 if temperature is None else temperature, + token_length_normalized if choices_method is None else choices_method, + ) + + # check regex is valid + if regex is not None: + try: + re.compile(regex) + except re.error as e: + raise e + + return SglGen( + name, + max_tokens, + min_tokens, + n, + stop, + stop_token_ids, + temperature, + top_p, + top_k, + min_p, + frequency_penalty, + presence_penalty, + ignore_eos, + return_logprob, + logprob_start_len, + top_logprobs_num, + return_text_in_logprobs, + dtype, + regex, + json_schema, + ) + + +def gen_int( + name: Optional[str] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, +): + return SglGen( + name, + max_tokens, + None, + n, + stop, + stop_token_ids, + temperature, + top_p, + top_k, + min_p, + frequency_penalty, + presence_penalty, + ignore_eos, + return_logprob, + logprob_start_len, + top_logprobs_num, + return_text_in_logprobs, + int, + None, + ) + + +def gen_string( + name: Optional[str] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, +): + return SglGen( + name, + max_tokens, + None, + n, + stop, + stop_token_ids, + temperature, + top_p, + top_k, + min_p, + frequency_penalty, + presence_penalty, + ignore_eos, + return_logprob, + logprob_start_len, + top_logprobs_num, + return_text_in_logprobs, + str, + None, + ) + + +def image(expr: SglExpr): + return SglImage(expr) + + +def video(path: str, num_frames: int): + return SglVideo(path, num_frames) + + +def select( + name: Optional[str] = None, + choices: Optional[List[str]] = None, + temperature: float = 0.0, + choices_method: ChoicesSamplingMethod = token_length_normalized, +): + assert choices is not None + return SglSelect(name, choices, temperature, choices_method) + + +def _role_common(name: str, expr: Optional[SglExpr] = None): + if expr is None: + return SglExprList([SglRoleBegin(name), SglRoleEnd(name)]) + else: + return SglExprList([SglRoleBegin(name), expr, SglRoleEnd(name)]) + + +def system(expr: Optional[SglExpr] = None): + return _role_common("system", expr) + + +def user(expr: Optional[SglExpr] = None): + return _role_common("user", expr) + + +def assistant(expr: Optional[SglExpr] = None): + return _role_common("assistant", expr) + + +def system_begin(): + return SglRoleBegin("system") + + +def system_end(): + return SglRoleEnd("system") + + +def user_begin(): + return SglRoleBegin("user") + + +def user_end(): + return SglRoleEnd("user") + + +def assistant_begin(): + return SglRoleBegin("assistant") + + +def assistant_end(): + return SglRoleEnd("assistant") + + +def separate_reasoning( + expr: Optional[SglExpr] = None, model_type: Optional[str] = None +): + return SglExprList([expr, SglSeparateReasoning(model_type, expr=expr)]) diff --git a/python/sglang/lang/backend/anthropic.py b/python/sglang/lang/backend/anthropic.py new file mode 100644 index 000000000..4918a1703 --- /dev/null +++ b/python/sglang/lang/backend/anthropic.py @@ -0,0 +1,73 @@ +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.chat_template import get_chat_template +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + +try: + import anthropic +except ImportError as e: + anthropic = e + + +class Anthropic(BaseBackend): + def __init__(self, model_name, *args, **kwargs): + super().__init__() + + if isinstance(anthropic, Exception): + raise anthropic + + self.model_name = model_name + self.chat_template = get_chat_template("claude") + self.client = anthropic.Anthropic(*args, **kwargs) + + def get_chat_template(self): + return self.chat_template + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + messages = s.messages_ + else: + messages = [{"role": "user", "content": s.text_}] + + if messages and messages[0]["role"] == "system": + system = messages.pop(0)["content"] + else: + system = "" + + ret = self.client.messages.create( + model=self.model_name, + system=system, + messages=messages, + **sampling_params.to_anthropic_kwargs(), + ) + comp = ret.content[0].text + + return comp, {} + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + messages = s.messages_ + else: + messages = [{"role": "user", "content": s.text_}] + + if messages and messages[0]["role"] == "system": + system = messages.pop(0)["content"] + else: + system = "" + + with self.client.messages.stream( + model=self.model_name, + system=system, + messages=messages, + **sampling_params.to_anthropic_kwargs(), + ) as stream: + for text in stream.text_stream: + yield text, {} diff --git a/python/sglang/lang/backend/base_backend.py b/python/sglang/lang/backend/base_backend.py new file mode 100644 index 000000000..62dd50416 --- /dev/null +++ b/python/sglang/lang/backend/base_backend.py @@ -0,0 +1,82 @@ +from typing import List, Optional, Union + +from sglang.lang.chat_template import get_chat_template +from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + + +class BaseBackend: + def __init__(self) -> None: + self.support_concate_and_append = False + self.chat_template = get_chat_template("default") + + def get_model_name(self): + raise NotImplementedError() + + def get_chat_template(self): + return self.chat_template + + def cache_prefix(self, prefix_str: str): + pass + + def uncache_prefix(self, rid: str): + pass + + def end_request(self, rid: Union[str, List[str]]): + pass + + def begin_program(self, s: StreamExecutor): + pass + + def end_program(self, s: Union[StreamExecutor, List[StreamExecutor]]): + pass + + def commit_lazy_operations(self, s: StreamExecutor): + pass + + def fork_program( + self, + src: StreamExecutor, + dst: List[StreamExecutor], + position_ids_offset: Optional[List[int]] = None, + ): + pass + + def fill_image(self, s: StreamExecutor): + pass + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + raise NotImplementedError() + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + raise NotImplementedError() + + def select( + self, + s: StreamExecutor, + choices: List[str], + temperature: float, + choices_method: Optional[ChoicesSamplingMethod] = None, + ) -> ChoicesDecision: + raise NotImplementedError() + + def concatenate_and_append(self, src_rids: List[str], dst_rid: str): + raise NotImplementedError() + + def shutdown(self): + pass + + def flush_cache(self): + pass + + def get_server_info(self): + pass diff --git a/python/sglang/lang/backend/litellm.py b/python/sglang/lang/backend/litellm.py new file mode 100644 index 000000000..5803b5431 --- /dev/null +++ b/python/sglang/lang/backend/litellm.py @@ -0,0 +1,90 @@ +from typing import Mapping, Optional + +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.chat_template import get_chat_template_by_model_path +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + +try: + import litellm +except ImportError as e: + litellm = e + litellm.num_retries = 1 + + +class LiteLLM(BaseBackend): + def __init__( + self, + model_name, + chat_template=None, + api_key=None, + organization: Optional[str] = None, + base_url: Optional[str] = None, + timeout: Optional[float] = 600, + max_retries: Optional[int] = litellm.num_retries, + default_headers: Optional[Mapping[str, str]] = None, + ): + super().__init__() + + if isinstance(litellm, Exception): + raise litellm + + self.model_name = model_name + + self.chat_template = chat_template or get_chat_template_by_model_path( + model_name + ) + + self.client_params = { + "api_key": api_key, + "organization": organization, + "base_url": base_url, + "timeout": timeout, + "max_retries": max_retries, + "default_headers": default_headers, + } + + def get_chat_template(self): + return self.chat_template + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + messages = s.messages_ + else: + messages = [{"role": "user", "content": s.text_}] + + ret = litellm.completion( + model=self.model_name, + messages=messages, + **self.client_params, + **sampling_params.to_litellm_kwargs(), + ) + comp = ret.choices[0].message.content + + return comp, {} + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + messages = s.messages_ + else: + messages = [{"role": "user", "content": s.text_}] + + ret = litellm.completion( + model=self.model_name, + messages=messages, + stream=True, + **self.client_params, + **sampling_params.to_litellm_kwargs(), + ) + for chunk in ret: + text = chunk.choices[0].delta.content + if text is not None: + yield text, {} diff --git a/python/sglang/lang/backend/openai.py b/python/sglang/lang/backend/openai.py new file mode 100644 index 000000000..a2d006bb7 --- /dev/null +++ b/python/sglang/lang/backend/openai.py @@ -0,0 +1,475 @@ +import dataclasses +import logging +import time +import warnings +from typing import List, Optional, Union + +import numpy as np + +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.chat_template import ChatTemplate, get_chat_template_by_model_path +from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + +try: + import openai + import tiktoken +except ImportError as e: + openai = tiktoken = e + + +logger = logging.getLogger(__name__) + + +def create_logit_bias_int(tokenizer): + """Get logit bias for integer numbers.""" + int_token_ids = [] + + tokens = tokenizer._mergeable_ranks + for token, token_id in tokens.items(): + s = tokenizer.decode([token_id]) + if all([c.isdigit() for c in s]) or s in [" "]: + int_token_ids.append(token_id) + if len(int_token_ids) >= 300: # OpenAI API limit + break + special_tokens = tokenizer._special_tokens + mask = {t: 100 for t in int_token_ids[:299]} + mask[special_tokens["<|endoftext|>"]] = 100 + return mask + + +INSTRUCT_MODEL_NAMES = [ + "gpt-3.5-turbo-instruct", +] + + +@dataclasses.dataclass +class TokenUsage: + prompt_tokens: int + completion_tokens: int + + def reset(self): + self.prompt_tokens = self.completion_tokens = 0 + + +class OpenAI(BaseBackend): + def __init__( + self, + model_name: str, + is_chat_model: Optional[bool] = None, + chat_template: Optional[ChatTemplate] = None, + is_azure: bool = False, + *args, + **kwargs, + ): + super().__init__() + + if isinstance(openai, Exception): + raise openai + + if is_azure: + self.client = openai.AzureOpenAI(*args, **kwargs) + else: + self.client = openai.OpenAI(*args, **kwargs) + + self.model_name = model_name + try: + self.tokenizer = tiktoken.encoding_for_model(model_name) + except KeyError: + self.tokenizer = tiktoken.get_encoding("cl100k_base") + self.logit_bias_int = create_logit_bias_int(self.tokenizer) + + self.chat_template = chat_template or get_chat_template_by_model_path( + model_name + ) + + if is_chat_model is not None: + self.is_chat_model = is_chat_model + else: + if model_name in INSTRUCT_MODEL_NAMES: + self.is_chat_model = False + else: + self.is_chat_model = True + + self.chat_prefix = self.chat_template.role_prefix_and_suffix["assistant"][0] + + # Usage + self.token_usage = TokenUsage(0, 0) + + # API speculative execution + # TODO(ying): This does not support multi-threading (run_batch) + self.spec_kwargs = {} + self.spec_format = [] + self.spec_max_num_tries = 3 + + def get_chat_template(self): + return self.chat_template + + def _prepare_spec_execution( + self, + sampling_params: SglSamplingParams, + num_api_spec_tokens: int, + spec_var_name: str, + ): + if "max_tokens" not in self.spec_kwargs: + self.spec_kwargs["max_tokens"] = num_api_spec_tokens + else: + assert self.spec_kwargs["max_tokens"] == num_api_spec_tokens + + params = sampling_params.to_openai_kwargs() + for key, value in params.items(): + if key in ["stop"]: + continue + if key in ["max_tokens"]: + warnings.warn( + "The parameter max_tokens will be overwritten by speculated number of tokens." + ) + continue + if key not in self.spec_kwargs: + self.spec_kwargs[key] = value + else: + assert ( + value == self.spec_kwargs[key] + ), "sampling parameters should be consistent if turn on api speculative execution." + self.spec_format.append( + {"text": "", "stop": params["stop"], "name": spec_var_name} + ) + return "", {} + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + spec_var_name: str = None, + ): + if sampling_params.dtype is None: + if self.is_chat_model: + if s.num_api_spec_tokens is None: + if not s.text_.endswith(self.chat_prefix): + raise RuntimeError( + "This use case is not supported if api speculative execution is off. " + "For OpenAI chat models, sgl.gen must be right after sgl.assistant. " + "Example of adding api speculative execution: @function(num_api_spec_tokens=128)." + ) + prompt = s.messages_ + else: + return self._prepare_spec_execution( + sampling_params, s.num_api_spec_tokens, spec_var_name + ) + else: + prompt = s.text_ + + kwargs = sampling_params.to_openai_kwargs() + if ( + self.model_name.startswith("o1") + or self.model_name.startswith("o3") + or "o1" in self.model_name + ): + kwargs.pop("max_tokens", None) + else: + kwargs.pop("max_completion_tokens", None) + + comp = openai_completion( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=prompt, + **kwargs, + ) + # Keep the returned list (or string) as is. + elif sampling_params.dtype in [str, "str", "string"]: + assert ( + not self.is_chat_model + ), "constrained type not supported on chat model" + kwargs = sampling_params.to_openai_kwargs() + kwargs.pop("stop") + comp = openai_completion( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=s.text_ + '"', + stop='"', + **kwargs, + ) + # Wrap each element in quotes if we have a list. + if isinstance(comp, list): + comp = ['"' + x + '"' for x in comp] + else: + comp = '"' + comp + '"' + elif sampling_params.dtype in [int, "int"]: + assert ( + not self.is_chat_model + ), "constrained type not supported on chat model" + kwargs = sampling_params.to_openai_kwargs() + kwargs.pop("stop") + comp = openai_completion( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=s.text_, + logit_bias=self.logit_bias_int, + stop=[" "], + **kwargs, + ) + # Leave as a list if that's what is returned. + else: + raise ValueError(f"Unknown dtype: {sampling_params.dtype}") + + return comp, {} + + def spec_fill(self, value: str): + assert self.is_chat_model + self.spec_format.append({"text": value, "stop": None, "name": None}) + + def spec_pattern_match(self, comp): + for i, term in enumerate(self.spec_format): + text = term["text"] + if text != "": + if comp.startswith(text): + comp = comp[len(text) :] + else: + return False + else: + pos = comp.find(term["stop"]) + if pos != -1: + term["text"] = comp[:pos] + comp = comp[pos:] + else: + if i == len(self.spec_format) - 1: + term["text"] = comp + else: + return False + return True + + def role_end_generate( + self, + s: StreamExecutor, + ): + if s.num_api_spec_tokens is None or not s.text_.endswith(self.chat_prefix): + return + + comp = "" + if not all(x["name"] is None for x in self.spec_format): + # TODO(ying): throw errors or warnings + for i in range(self.spec_max_num_tries): + comp = openai_completion( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=s.messages_, + **self.spec_kwargs, + ) + # Use a string for pattern matching. + comp_for_match = comp[0] if isinstance(comp, list) else comp + if self.spec_pattern_match(comp_for_match): + break + + for term in self.spec_format: + s.text_ += term["text"] + name = term["name"] + if name is not None: + s.variables[name] = term["text"] + s.meta_info[name] = {} + s.variable_event[name].set() + + self.spec_kwargs = {} + self.spec_format = [] + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if sampling_params.dtype is None: + if self.is_chat_model: + if not s.text_.endswith(self.chat_prefix): + raise RuntimeError( + "This use case is not supported. " + "For OpenAI chat models, sgl.gen must be right after sgl.assistant" + ) + prompt = s.messages_ + else: + prompt = s.text_ + + kwargs = sampling_params.to_openai_kwargs() + generator = openai_completion_stream( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=prompt, + **kwargs, + ) + return generator + else: + raise ValueError(f"Unknown dtype: {sampling_params.dtype}") + + def select( + self, + s: StreamExecutor, + choices: List[str], + temperature: float, + choices_method: ChoicesSamplingMethod, + ) -> ChoicesDecision: + """Note: `choices_method` is not used by the OpenAI backend.""" + if self.is_chat_model: + raise NotImplementedError( + "select/choices is not supported for chat models. " + "Please try to use a non-chat model such as gpt-3.5-turbo-instruct" + ) + + n_choices = len(choices) + token_ids = [self.tokenizer.encode(x) for x in choices] + scores = [0] * n_choices + valid = [len(x) > 0 for x in token_ids] + prompt_tokens = self.tokenizer.encode(s.text_) + + max_len = max([len(x) for x in token_ids]) + for step in range(max_len): + # Build logit bias + logit_bias = {} + for i in range(n_choices): + if valid[i]: + logit_bias[token_ids[i][step]] = 100 + + # Call API + ret = self.client.completions.create( + model=self.model_name, + prompt=prompt_tokens, + logit_bias=logit_bias, + max_tokens=1, + temperature=temperature, + ) + ret_str = ret.choices[0].text + ret_token = self.tokenizer.encode(ret_str)[0] + self.token_usage.prompt_tokens += ret.usage.prompt_tokens + self.token_usage.completion_tokens = ret.usage.completion_tokens + + # TODO: + # 1. return logits as the scores + # 2. compute logits of the full choice + # 3. consider chunk-based decoding + + # Update valid + hit = False + for i in range(n_choices): + if valid[i]: + if step == len(token_ids[i]) - 1: + valid[i] = False + + if ret_token == token_ids[i][step]: + scores[i] += 1 + hit = True + else: + valid[i] = False + assert hit + + if np.sum(valid) <= 1: + break + + prompt_tokens.append(ret_token) + + return ChoicesDecision( + decision=choices[np.argmax(scores)], + meta_info={"scores": scores}, + ) + + +def openai_completion( + client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs +) -> Union[str, List[str]]: + # if "ebnf" is in kwargs, warn and remove + if "ebnf" in kwargs: + warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.") + del kwargs["ebnf"] + + for attempt in range(retries): + try: + if is_chat: + if "stop" in kwargs and kwargs["stop"] is None: + kwargs.pop("stop") + ret = client.chat.completions.create(messages=prompt, **kwargs) + if len(ret.choices) == 1: + comp = ret.choices[0].message.content + else: + comp = [c.message.content for c in ret.choices] + else: + ret = client.completions.create(prompt=prompt, **kwargs) + if isinstance(prompt, (list, tuple)): + comp = [c.text for c in ret.choices] + else: + comp = ret.choices[0].text + if len(ret.choices) > 1: + comp = [c.text for c in ret.choices] + + token_usage.prompt_tokens += ret.usage.prompt_tokens + token_usage.completion_tokens += ret.usage.completion_tokens + break + except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e: + logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...") + time.sleep(5) + if attempt == retries - 1: + raise e + except Exception as e: + logger.error(f"RuntimeError {e}.") + raise e + + return comp + + +def openai_completion_stream( + client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs +): + # if "ebnf" is in kwargs, warn and remove + if "ebnf" in kwargs: + warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.") + del kwargs["ebnf"] + + for attempt in range(retries): + try: + if is_chat: + if "stop" in kwargs and kwargs["stop"] is None: + kwargs.pop("stop") + generator = client.chat.completions.create( + messages=prompt, + stream=True, + stream_options={"include_usage": True}, + **kwargs, + ) + for ret in generator: + if len(ret.choices) == 0: + continue + try: + content = ret.choices[0].delta.content + except IndexError: + content = None + yield content or "", {} + else: + generator = client.completions.create( + prompt=prompt, + stream=True, + stream_options={"include_usage": True}, + **kwargs, + ) + for ret in generator: + if len(ret.choices) == 0: + continue + content = ret.choices[0].text + yield content or "", {} + + token_usage.prompt_tokens += ret.usage.prompt_tokens + token_usage.completion_tokens += ret.usage.completion_tokens + break + except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e: + logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...") + time.sleep(5) + if attempt == retries - 1: + raise e + except Exception as e: + logger.error(f"RuntimeError {e}.") + raise e diff --git a/python/sglang/lang/backend/runtime_endpoint.py b/python/sglang/lang/backend/runtime_endpoint.py new file mode 100644 index 000000000..349f9934a --- /dev/null +++ b/python/sglang/lang/backend/runtime_endpoint.py @@ -0,0 +1,527 @@ +import atexit +import json +import multiprocessing +import warnings +from typing import Dict, List, Optional, Union + +import aiohttp +import requests + +from sglang.global_config import global_config +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.chat_template import get_chat_template, get_chat_template_by_model_path +from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import ( + REGEX_BOOL, + REGEX_FLOAT, + REGEX_INT, + REGEX_STR, + SglSamplingParams, +) +from sglang.utils import http_request + + +class RuntimeEndpoint(BaseBackend): + def __init__( + self, + base_url: str, + api_key: Optional[str] = None, + verify: Optional[str] = None, + chat_template_name: Optional[str] = None, + ): + super().__init__() + self.support_concate_and_append = True + + self.base_url = base_url + self.api_key = api_key + self.verify = verify + + res = http_request( + self.base_url + "/get_model_info", + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + self.model_info = res.json() + + if chat_template_name: + self.chat_template = get_chat_template(chat_template_name) + else: + self.chat_template = get_chat_template_by_model_path( + self.model_info["model_path"] + ) + + def get_model_name(self): + return self.model_info["model_path"] + + def flush_cache(self): + res = http_request( + self.base_url + "/flush_cache", + api_key=self.api_key, + verify=self.verify, + method="POST", + ) + self._assert_success(res) + + def get_server_info(self): + res = http_request( + self.base_url + "/get_server_info", + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + return res.json() + + def get_chat_template(self): + return self.chat_template + + def cache_prefix(self, prefix_str: str): + res = http_request( + self.base_url + "/generate", + json={"text": prefix_str, "sampling_params": {"max_new_tokens": 0}}, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def start_profile(self): + res = http_request( + self.base_url + "/start_profile", + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def stop_profile(self): + res = http_request( + self.base_url + "/stop_profile", + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def commit_lazy_operations(self, s: StreamExecutor): + data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}} + self._add_images(s, data) + res = http_request( + self.base_url + "/generate", + json=data, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def fill_image(self, s: StreamExecutor): + data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}} + self._add_images(s, data) + res = http_request( + self.base_url + "/generate", + json=data, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def _handle_dtype_to_regex(self, sampling_params: SglSamplingParams): + if sampling_params.dtype is None: + return + + if sampling_params.stop == (): + sampling_params.stop = [] + + dtype_regex = None + if sampling_params.dtype in ["int", int]: + + dtype_regex = REGEX_INT + sampling_params.stop.extend([" ", "\n"]) + elif sampling_params.dtype in ["float", float]: + + dtype_regex = REGEX_FLOAT + sampling_params.stop.extend([" ", "\n"]) + elif sampling_params.dtype in ["str", str]: + + dtype_regex = REGEX_STR + elif sampling_params.dtype in ["bool", bool]: + + dtype_regex = REGEX_BOOL + else: + raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}") + + if dtype_regex is not None and sampling_params.regex is not None: + warnings.warn( + f"Both dtype and regex are set. Only dtype will be used. dtype: {sampling_params.dtype}, regex: {sampling_params.regex}" + ) + + sampling_params.regex = dtype_regex + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + self._handle_dtype_to_regex(sampling_params) + data = { + "text": s.text_, + "sampling_params": { + "skip_special_tokens": global_config.skip_special_tokens_in_output, + "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out, + **sampling_params.to_srt_kwargs(), + }, + } + + for item in [ + "return_logprob", + "logprob_start_len", + "top_logprobs_num", + "return_text_in_logprobs", + ]: + value = getattr(sampling_params, item, None) + if value is not None: + data[item] = value + + self._add_images(s, data) + + res = http_request( + self.base_url + "/generate", + json=data, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + obj = res.json() + comp = obj["text"] + return comp, obj["meta_info"] + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + self._handle_dtype_to_regex(sampling_params) + + data = { + "text": s.text_, + "sampling_params": { + "skip_special_tokens": global_config.skip_special_tokens_in_output, + "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out, + **sampling_params.to_srt_kwargs(), + }, + } + + for item in [ + "return_logprob", + "logprob_start_len", + "top_logprobs_num", + "return_text_in_logprobs", + ]: + value = getattr(sampling_params, item, None) + if value is not None: + data[item] = value + + data["stream"] = True + self._add_images(s, data) + + res = http_request( + self.base_url + "/generate", + json=data, + stream=True, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + pos = 0 + + for chunk in res.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + chunk_text = data["text"][pos:] + meta_info = data["meta_info"] + pos += len(chunk_text) + yield chunk_text, meta_info + + def select( + self, + s: StreamExecutor, + choices: List[str], + temperature: float, + choices_method: ChoicesSamplingMethod, + ) -> ChoicesDecision: + assert temperature <= 1e-5 + + # Cache common prefix + data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}} + obj = self._generate_http_request(s, data) + prompt_len = obj["meta_info"]["prompt_tokens"] + logprob_start_len = max(prompt_len - 2, 0) # For token healing + + # Compute logprob + data = { + "text": [s.text_ + c for c in choices], + "sampling_params": { + "max_new_tokens": 0, + "temperature": 0, + }, + "return_logprob": True, + "return_text_in_logprobs": True, + "logprob_start_len": logprob_start_len, + } + obj = self._generate_http_request(s, data) + + input_token_logprobs = [r["meta_info"]["input_token_logprobs"] for r in obj] + output_token_logprobs = [r["meta_info"]["output_token_logprobs"] for r in obj] + normalized_prompt_logprobs = [ + compute_normalized_prompt_logprobs(r["meta_info"]["input_token_logprobs"]) + for r in obj + ] + + # Remove extra token if no token healing occurred + for i in range(len(input_token_logprobs)): + healed_token_str = input_token_logprobs[i][0][-1] + if s.text_.endswith(healed_token_str): + healed_token_logprob = input_token_logprobs[i][0][0] + normalized_prompt_logprobs[i] = ( + normalized_prompt_logprobs[i] * len(input_token_logprobs[i]) + - healed_token_logprob + ) / (len(input_token_logprobs[i]) - 1) + input_token_logprobs[i] = input_token_logprobs[i][1:] + + # Compute unconditional logprobs if required + if choices_method.requires_unconditional_logprobs: + input_ids = [[el[1] for el in subl] for subl in input_token_logprobs] + data = { + "input_ids": input_ids, + "sampling_params": {"max_new_tokens": 0}, + "return_logprob": True, + } + obj = self._generate_http_request(s, data) + unconditional_token_logprobs = [ + r["meta_info"]["input_token_logprobs"] for r in obj + ] + else: + unconditional_token_logprobs = None + + return choices_method( + choices=choices, + normalized_prompt_logprobs=normalized_prompt_logprobs, + input_token_logprobs=input_token_logprobs, + output_token_logprobs=output_token_logprobs, + unconditional_token_logprobs=unconditional_token_logprobs, + ) + + def concatenate_and_append(self, src_rids: List[str], dst_rid: str): + res = http_request( + self.base_url + "/concate_and_append_request", + json={"src_rids": src_rids, "dst_rid": dst_rid}, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def _generate_http_request(self, s: StreamExecutor, data): + self._add_images(s, data) + res = http_request( + self.base_url + "/generate", + json=data, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + return res.json() + + def _add_images(self, s: StreamExecutor, data): + if s.images_: + assert len(s.images_) == 1, "Only support one image." + data["image_data"] = s.images_[0][1] + + def _assert_success(self, res): + if res.status_code != 200: + try: + content = res.json() + except json.JSONDecodeError: + content = res.text + raise RuntimeError(content) + + +def compute_normalized_prompt_logprobs(input_logprobs): + values = [x[0] for x in input_logprobs if x[0]] + return sum(values) / len(values) + + +class Runtime: + """ + A wrapper for the HTTP server. + This is used for launching the server in a python program without + using the command line interface. + + It is mainly used for the frontend language. + You should use the Engine class if you want to do normal offline processing without the frontend language. + """ + + def __init__( + self, + log_level: str = "error", + *args, + **kwargs, + ): + """See the arguments in server_args.py::ServerArgs""" + # We delay the import of any `sglang.srt` components in `sglang.lang`, so users can run + # client code without installing SRT server and its dependency if they want. + from sglang.srt.entrypoints.http_server import launch_server + from sglang.srt.server_args import ServerArgs + from sglang.srt.utils import is_port_available + + self.server_args = ServerArgs(*args, log_level=log_level, **kwargs) + + # Pre-allocate ports + for port in range(self.server_args.port, 40000): + if is_port_available(port): + break + self.server_args.port = port + + self.url = self.server_args.url() + self.generate_url = self.url + "/generate" + + # NOTE: We store pid instead of proc to fix some issues during __delete__ + self.pid = None + pipe_reader, pipe_writer = multiprocessing.Pipe(duplex=False) + + ctx = multiprocessing.get_context("spawn") + proc = ctx.Process( + target=launch_server, + args=(self.server_args, pipe_writer), + ) + proc.start() + pipe_writer.close() + self.pid = proc.pid + + # Before python program terminates, call shutdown implicitly. Therefore, users don't have to explicitly call .shutdown() + atexit.register(self.shutdown) + + # TODO: remove this pipe_writer mechanism and use `/health_generate` instead. + try: + init_state = pipe_reader.recv() + except EOFError: + init_state = "" + + if init_state != "ready": + self.shutdown() + raise RuntimeError( + "Initialization failed. Please see the error messages above." + ) + + self.endpoint = RuntimeEndpoint(self.url) + + def shutdown(self): + from sglang.srt.utils import kill_process_tree + + if self.pid is not None: + kill_process_tree(self.pid) + self.pid = None + + def start_profile(self): + self.endpoint.start_profile() + + def stop_profile(self): + self.endpoint.stop_profile() + + def cache_prefix(self, prefix: str): + self.endpoint.cache_prefix(prefix) + + def get_tokenizer(self): + from sglang.srt.hf_transformers_utils import get_tokenizer + + return get_tokenizer( + self.server_args.tokenizer_path, + tokenizer_mode=self.server_args.tokenizer_mode, + trust_remote_code=self.server_args.trust_remote_code, + revision=self.server_args.revision, + ) + + async def async_generate( + self, + prompt: str, + sampling_params: Optional[Dict] = None, + ): + if self.server_args.skip_tokenizer_init: + json_data = { + "input_ids": prompt, + "sampling_params": sampling_params, + "stream": True, + } + else: + json_data = { + "text": prompt, + "sampling_params": sampling_params, + "stream": True, + } + pos = 0 + + timeout = aiohttp.ClientTimeout(total=3 * 3600) + async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session: + async with session.post(self.generate_url, json=json_data) as response: + async for chunk, _ in response.content.iter_chunks(): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]\n\n": + break + data = json.loads(chunk[5:].strip("\n")) + if "text" in data: + cur = data["text"][pos:] + if cur: + yield cur + pos += len(cur) + else: + yield data + + add_request = async_generate + + def generate( + self, + prompt: Union[str, List[str]], + sampling_params: Optional[Dict] = None, + return_logprob: Optional[Union[List[bool], bool]] = False, + logprob_start_len: Optional[Union[List[int], int]] = None, + top_logprobs_num: Optional[Union[List[int], int]] = None, + lora_path: Optional[List[Optional[str]]] = None, + ): + json_data = { + "text": prompt, + "sampling_params": sampling_params, + "return_logprob": return_logprob, + "logprob_start_len": logprob_start_len, + "top_logprobs_num": top_logprobs_num, + "lora_path": lora_path, + } + assert not isinstance(lora_path, list) or len(lora_path) == len(prompt) + response = requests.post( + self.url + "/generate", + json=json_data, + ) + return json.dumps(response.json()) + + def encode( + self, + prompt: Union[str, List[str], List[Dict], List[List[Dict]]], + ): + json_data = {"text": prompt} + response = requests.post(self.url + "/encode", json=json_data) + return json.dumps(response.json()) + + async def get_server_info(self): + async with aiohttp.ClientSession() as session: + async with session.get(f"{self.url}/get_server_info") as response: + if response.status == 200: + return await response.json() + else: + error_data = await response.json() + raise RuntimeError( + f"Failed to get server info. {error_data['error']['message']}" + ) + + def __del__(self): + self.shutdown() diff --git a/python/sglang/lang/backend/vertexai.py b/python/sglang/lang/backend/vertexai.py new file mode 100644 index 000000000..3d51fb137 --- /dev/null +++ b/python/sglang/lang/backend/vertexai.py @@ -0,0 +1,148 @@ +import os +import warnings + +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.chat_template import get_chat_template +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + +try: + import vertexai + from vertexai.preview.generative_models import ( + GenerationConfig, + GenerativeModel, + Image, + ) +except ImportError as e: + GenerativeModel = e + + +class VertexAI(BaseBackend): + def __init__(self, model_name, safety_settings=None): + super().__init__() + + if isinstance(GenerativeModel, Exception): + raise GenerativeModel + + project_id = os.environ["GCP_PROJECT_ID"] + location = os.environ.get("GCP_LOCATION") + vertexai.init(project=project_id, location=location) + + self.model_name = model_name + self.chat_template = get_chat_template("default") + self.safety_settings = safety_settings + + def get_chat_template(self): + return self.chat_template + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + prompt = self.messages_to_vertexai_input(s.messages_) + else: + # single-turn + prompt = ( + self.text_to_vertexai_input(s.text_, s.cur_images) + if s.cur_images + else s.text_ + ) + ret = GenerativeModel(self.model_name).generate_content( + prompt, + generation_config=GenerationConfig(**sampling_params.to_vertexai_kwargs()), + safety_settings=self.safety_settings, + ) + + comp = ret.text + + return comp, {} + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + prompt = self.messages_to_vertexai_input(s.messages_) + else: + # single-turn + prompt = ( + self.text_to_vertexai_input(s.text_, s.cur_images) + if s.cur_images + else s.text_ + ) + generator = GenerativeModel(self.model_name).generate_content( + prompt, + stream=True, + generation_config=GenerationConfig(**sampling_params.to_vertexai_kwargs()), + safety_settings=self.safety_settings, + ) + for ret in generator: + yield ret.text, {} + + def text_to_vertexai_input(self, text, images): + input = [] + # split with image token + text_segs = text.split(self.chat_template.image_token) + for image_path, image_base64_data in images: + text_seg = text_segs.pop(0) + if text_seg != "": + input.append(text_seg) + input.append(Image.from_bytes(image_base64_data)) + text_seg = text_segs.pop(0) + if text_seg != "": + input.append(text_seg) + return input + + def messages_to_vertexai_input(self, messages): + vertexai_message = [] + # from openai message format to vertexai message format + for msg in messages: + if isinstance(msg["content"], str): + text = msg["content"] + else: + text = msg["content"][0]["text"] + + if msg["role"] == "system": + warnings.warn("Warning: system prompt is not supported in VertexAI.") + vertexai_message.append( + { + "role": "user", + "parts": [{"text": "System prompt: " + text}], + } + ) + vertexai_message.append( + { + "role": "model", + "parts": [{"text": "Understood."}], + } + ) + continue + if msg["role"] == "user": + vertexai_msg = { + "role": "user", + "parts": [{"text": text}], + } + elif msg["role"] == "assistant": + vertexai_msg = { + "role": "model", + "parts": [{"text": text}], + } + + # images + if isinstance(msg["content"], list) and len(msg["content"]) > 1: + for image in msg["content"][1:]: + assert image["type"] == "image_url" + vertexai_msg["parts"].append( + { + "inline_data": { + "data": image["image_url"]["url"].split(",")[1], + "mime_type": "image/jpeg", + } + } + ) + + vertexai_message.append(vertexai_msg) + return vertexai_message diff --git a/python/sglang/lang/chat_template.py b/python/sglang/lang/chat_template.py new file mode 100644 index 000000000..80ea6d963 --- /dev/null +++ b/python/sglang/lang/chat_template.py @@ -0,0 +1,662 @@ +import re +from dataclasses import dataclass +from enum import Enum, auto +from typing import Callable, Dict, List, Tuple + + +class ChatTemplateStyle(Enum): + PLAIN = auto() + LLAMA2 = auto() + + +@dataclass +class ChatTemplate: + name: str + default_system_prompt: str + role_prefix_and_suffix: Dict[str, Tuple[str, str]] + stop_str: List[str] = () + image_token: str = "" + audio_token: str = ""), + }, + image_token=" \n", + ) +) + +register_chat_template( + ChatTemplate( + name="llama-2-chat", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("<>\n", "\n<>\n\n"), + "user": ("[INST] ", " [/INST]"), + "assistant": ("", " "), + }, + style=ChatTemplateStyle.LLAMA2, + ) +) + +# Reference: https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503/blob/main/chat_template.json +register_chat_template( + ChatTemplate( + name="mistral", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("[SYSTEM_PROMPT] ", " [/SYSTEM_PROMPT]"), + "user": ("[INST] ", " [/INST]"), + "assistant": ("", " "), + }, + stop_str=("",), + image_token="[IMG]", + ) +) + +register_chat_template( + ChatTemplate( + name="llama-3-instruct", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "<|start_header_id|>system<|end_header_id|>\n\n", + "<|eot_id|>", + ), + "user": ( + "<|start_header_id|>user<|end_header_id|>\n\n", + "<|eot_id|>", + ), + "assistant": ( + "<|start_header_id|>assistant<|end_header_id|>\n\n", + "<|eot_id|>", + ), + }, + stop_str=("<|eot_id|>",), + image_token="<|image|>", + ) +) + +# https://huggingface.co/openbmb/MiniCPM-V-2_6 +register_chat_template( + ChatTemplate( + name="minicpmv", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("", " "), + "user": ("user:", " "), + "assistant": ("assistant:", ""), + }, + stop_str=("<|im_end|>", "<|endoftext|>"), + image_token="(./)", + ) +) + +register_chat_template( + ChatTemplate( + name="janus-pro", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "", + "", + ), + "User": ( + "<｜User｜>", + "", + ), + "assistant": ( + "<｜Assistant｜>", + "<｜end▁of▁sentence｜>", + ), + }, + stop_str=("<｜end▁of▁sentence｜>",), + image_token="\n", + ) +) + +# https://huggingface.co/openbmb/MiniCPM-o-2_6 +register_chat_template( + ChatTemplate( + name="minicpmo", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("", " "), + "user": ("user:", " "), + "assistant": ("assistant:", ""), + }, + stop_str=("<|im_end|>", "<|endoftext|>"), + image_token="(./)", + audio_token="()", + ) +) + +register_chat_template( + ChatTemplate( + name="janus", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "", + "", + ), + "user": ( + "<｜User｜>", + "", + ), + "assistant": ( + "<｜Assistant｜>", + "<｜end▁of▁sentence｜>", + ), + }, + stop_str=("<｜end▁of▁sentence｜>",), + image_token="\n", + ) +) + +# The difference between "llama-3-instruct-llava" and "llama-3-instruct" is that llava uses a different image_token. +register_chat_template( + ChatTemplate( + name="llama-3-instruct-llava", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "<|start_header_id|>system<|end_header_id|>\n\n", + "<|eot_id|>", + ), + "user": ( + "<|start_header_id|>user<|end_header_id|>\n\n", + "<|eot_id|>", + ), + "assistant": ( + "<|start_header_id|>assistant<|end_header_id|>\n\n", + "<|eot_id|>", + ), + }, + stop_str=("<|eot_id|>",), + image_token="\n", + ) +) + +# Reference: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct/blob/main/chat_template.json +register_chat_template( + ChatTemplate( + name="llama-4", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "<|header_start|>system<|header_end|>\n\n", + "<|eot|>", + ), + "user": ( + "<|header_start|>user<|header_end|>\n\n", + "<|eot|>", + ), + "assistant": ( + "<|header_start|>assistant<|header_end|>\n\n", + "<|eot|>", + ), + }, + stop_str=("<|eot|>",), + image_token="<|image|>", + ) +) + +# Reference: https://modelscope.cn/models/01ai/Yi-1.5-34B-Chat/file/view/master?fileName=tokenizer_config.json&status=1 +register_chat_template( + ChatTemplate( + name="yi-1.5", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("", ""), + "user": ("<|im_start|>user\n", "<|im_end|>\n<|im_start|>assistant\n"), + "assistant": ("", "<|im_end|>\n"), + }, + style=ChatTemplateStyle.PLAIN, + stop_str=("<|im_end|>",), + ) +) + +# Reference: https://github.com/01-ai/Yi/tree/main/VL#major-difference-with-llava +register_chat_template( + ChatTemplate( + name="yi-vl", + default_system_prompt=( + "This is a chat between an inquisitive human and an AI assistant. Assume the role of the AI assistant. Read all the images carefully, and respond to the human's questions with informative, helpful, detailed and polite answers." + "这是一个好奇的人类和一个人工智能助手之间的对话。假设你扮演这个AI助手的角色。仔细阅读所有的图像，并对人类的问题做出信息丰富、有帮助、详细的和礼貌的回答。" + ), + role_prefix_and_suffix={ + "system": ("", "\n\n"), + "user": ("### Human:", "\n"), + "assistant": ("### Assistant:", "\n"), + }, + image_token=" \n", + ) +) + +register_chat_template( + ChatTemplate( + name="gemma-it", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("", ""), + "user": ("user\n", "\n"), + "assistant": ("model\n", "\n"), + }, + style=ChatTemplateStyle.PLAIN, + ) +) + +register_chat_template( + ChatTemplate( + name="dbrx-instruct", + default_system_prompt="You are DBRX, created by Databricks. You were last updated in December 2023. You answer questions based on information available up to that point.\nYOU PROVIDE SHORT RESPONSES TO SHORT QUESTIONS OR STATEMENTS, but provide thorough responses to more complex and open-ended questions.\nYou assist with various tasks, from writing to coding (using markdown for code blocks — remember to use ``` with code, JSON, and tables).\n(You do not have real-time data access or code execution capabilities. You avoid stereotyping and provide balanced perspectives on controversial topics. You do not provide song lyrics, poems, or news articles and do not divulge details of your training data.)\nThis is your system prompt, guiding your responses. Do not reference it, just respond to the user. If you find yourself talking about this message, stop. You should be responding appropriately and usually that means not mentioning this.\nYOU DO NOT MENTION ANY OF THIS INFORMATION ABOUT YOURSELF UNLESS THE INFORMATION IS DIRECTLY PERTINENT TO THE USER'S QUERY.", + role_prefix_and_suffix={ + "system": ("<|im_start|>system\n", "<|im_end|>"), + "user": ("\n<|im_start|>user\n", "<|im_end|>"), + "assistant": ("\n<|im_start|>assistant\n", "<|im_end|>"), + }, + stop_str=("<|im_end|>",), + ) +) + +register_chat_template( + ChatTemplate( + name="c4ai-command-r", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>", + "<|END_OF_TURN_TOKEN|>", + ), + "user": ("<|START_OF_TURN_TOKEN|><|USER_TOKEN|>", "<|END_OF_TURN_TOKEN|>"), + "assistant": ( + "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", + "<|END_OF_TURN_TOKEN|>", + ), + }, + style=ChatTemplateStyle.PLAIN, + ) +) + +# Adapted from https://huggingface.co/OpenGVLab/InternVL2-4B/blob/main/modeling_intern_vit.py +register_chat_template( + ChatTemplate( + name="internvl-2-5", + default_system_prompt="你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。", + role_prefix_and_suffix={ + "system": ("<|im_start|>system\n", "<|im_end|>\n"), + "user": ("<|im_start|>user\n", "<|im_end|>\n"), + "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"), + }, + stop_str=["<|im_end|>", "<|action_end|>"], + ) +) + +register_chat_template( + ChatTemplate( + name="interns1", + default_system_prompt="You are an AI assistant whose name is Intern-S1 (书生大模型).\n- Intern-S1 (书生大模型) is a vision-language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n- Intern-S1 (书生大模型) can understand and communicate fluently in the language chosen by the user such as English and 中文.\nYou are an expert reasoner with extensive experience in all areas. You approach problems through systematic thinking and rigorous reasoning. Your response should reflect deep understanding and precise logical thinking, making your solution path and reasoning clear to others. Please put your thinking process within ... tags.", + role_prefix_and_suffix={ + "system": ("<|im_start|>system\n", "<|im_end|>\n"), + "user": ("<|im_start|>user\n", "<|im_end|>\n"), + "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"), + }, + stop_str=["<|im_end|>", "<|action_end|>"], + ) +) + +register_chat_template( + ChatTemplate( + name="granite-3-instruct", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "<|start_of_role|>system<|end_of_role|>", + "<|end_of_text|>", + ), + "user": ( + "<|start_of_role|>user<|end_of_role|>", + "<|end_of_text|>", + ), + "assistant": ( + "<|start_of_role|>assistant<|end_of_role|>", + "<|end_of_text|>", + ), + }, + stop_str=("<|end_of_text|>",), + ) +) + +register_chat_template( + ChatTemplate( + name="deepseek-v3", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "", + "", + ), + "user": ( + "<｜User｜>", + "", + ), + "assistant": ( + "<｜Assistant｜>", + "<｜end▁of▁sentence｜>", + ), + }, + stop_str=("<｜end▁of▁sentence｜>",), + ) +) + +# Reference: https://huggingface.co/docs/transformers/main/model_doc/glm4_v#usage-example +register_chat_template( + ChatTemplate( + name="glm-4v", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("<|system|>\n", "\n"), + "user": ("<|user|>\n", "\n"), + "assistant": ("<|assistant|>\n", "\n"), + }, + style=ChatTemplateStyle.PLAIN, + stop_str=["<|user|>", "<|endoftext|>", "<|observation|>"], + image_token="<|image|>", + ) +) + + +@register_chat_template_matching_function +def match_deepseek(model_path: str): + if re.search(r"deepseek-(v3|r1)", model_path, re.IGNORECASE) and not re.search( + r"base", model_path, re.IGNORECASE + ): + return "deepseek-v3" + + +@register_chat_template_matching_function +def match_deepseek_janus_pro(model_path: str): + if re.search(r"janus", model_path, re.IGNORECASE): + return "janus-pro" + + +@register_chat_template_matching_function +def match_dbrx(model_path: str): + if re.search(r"dbrx", model_path, re.IGNORECASE) and re.search( + r"instruct", model_path, re.IGNORECASE + ): + return "dbrx-instruct" + + +@register_chat_template_matching_function +def match_vicuna(model_path: str): + if re.search(r"vicuna|llava-v1\.5|llava-next-video-7b", model_path, re.IGNORECASE): + return "vicuna_v1.1" + + +@register_chat_template_matching_function +def match_llama2_chat(model_path: str): + if re.search( + r"llama-2.*chat|codellama.*instruct", + model_path, + re.IGNORECASE, + ): + return "llama-2-chat" + + +@register_chat_template_matching_function +def match_mistral(model_path: str): + if re.search(r"pixtral|(mistral|mixtral).*instruct", model_path, re.IGNORECASE): + return "mistral" + + +@register_chat_template_matching_function +def match_llama3_instruct(model_path: str): + if re.search(r"llama-3.*instruct", model_path, re.IGNORECASE): + return "llama-3-instruct" + + +@register_chat_template_matching_function +def match_chat_ml(model_path: str): + if re.search(r"tinyllama", model_path, re.IGNORECASE): + return "chatml" + if re.search(r"qwen.*vl", model_path, re.IGNORECASE): + return "qwen2-vl" + if re.search(r"glm[-_]?4(\.\d+)?v", model_path, re.IGNORECASE): + return "glm-4v" + if re.search(r"qwen.*(chat|instruct)", model_path, re.IGNORECASE) and not re.search( + r"llava", model_path, re.IGNORECASE + ): + return "qwen" + if re.search( + r"llava-v1\.6-34b|llava-v1\.6-yi-34b|llava-next-video-34b|llava-onevision-qwen2", + model_path, + re.IGNORECASE, + ): + return "chatml-llava" + + +@register_chat_template_matching_function +def match_chat_yi(model_path: str): + if re.search(r"yi-vl", model_path, re.IGNORECASE) and not re.search( + r"llava", model_path, re.IGNORECASE + ): + return "yi-vl" + elif re.search(r"yi-1\.5.*chat", model_path, re.IGNORECASE): + return "yi-1.5" + + +@register_chat_template_matching_function +def match_gemma_it(model_path: str): + if re.search(r"gemma.*it", model_path, re.IGNORECASE): + return "gemma-it" + + +@register_chat_template_matching_function +def match_openbmb_minicpm(model_path: str): + if re.search(r"minicpm-v", model_path, re.IGNORECASE): + return "minicpmv" + elif re.search(r"minicpm-o", model_path, re.IGNORECASE): + return "minicpmo" + + +@register_chat_template_matching_function +def match_c4ai_command_r(model_path: str): + if re.search(r"c4ai-command-r", model_path, re.IGNORECASE): + return "c4ai-command-r" + + +@register_chat_template_matching_function +def match_granite_instruct(model_path: str): + if re.search(r"granite.*instruct", model_path, re.IGNORECASE): + return "granite-3-instruct" + + +@register_chat_template_matching_function +def match_gemma3_instruct(model_path: str): + if re.search(r"gemma-3", model_path, re.IGNORECASE): + return "gemma-it" + + +@register_chat_template_matching_function +def match_internvl_chat(model_path: str): + if re.search(r"internvl2_5", model_path, re.IGNORECASE): + return "internvl-2-5" + + +@register_chat_template_matching_function +def match_interns1_chat(model_path: str): + if re.search(r"intern-s1", model_path, re.IGNORECASE): + return "interns1" + if re.search(r"interns1", model_path, re.IGNORECASE): + return "interns1" + + +if __name__ == "__main__": + messages = [ + {"role": "system", "content": None}, # None means default + # {"role": "system", "content": "You are a helpful, respectful and honest assistant."}, + {"role": "user", "content": "Hello!"}, + {"role": "assistant", "content": "Hi!"}, + {"role": "user", "content": "What can you do?"}, + {"role": "assistant", "content": "I can chat with you."}, + ] + + template = get_chat_template("llama-2-chat") + print(template.get_prompt(messages)) diff --git a/python/sglang/lang/choices.py b/python/sglang/lang/choices.py new file mode 100644 index 000000000..e52c6b362 --- /dev/null +++ b/python/sglang/lang/choices.py @@ -0,0 +1,164 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +import numpy as np + + +@dataclass +class ChoicesDecision: + decision: str + meta_info: Optional[Dict[str, Any]] = None + + +class ChoicesSamplingMethod(ABC): + + @property + def requires_unconditional_logprobs(self) -> bool: + return False + + @abstractmethod + def __call__( + self, + *, + choices: List[str], + normalized_prompt_logprobs: List[float], + input_token_logprobs: List[List[Any]], + output_token_logprobs: List[List[Any]], + unconditional_token_logprobs: Optional[List[List[Any]]] = None, + ) -> ChoicesDecision: ... + + +class TokenLengthNormalized(ChoicesSamplingMethod): + + def __call__( + self, + *, + choices: List[str], + normalized_prompt_logprobs: List[float], + input_token_logprobs: List[List[Any]], + output_token_logprobs: List[List[Any]], + unconditional_token_logprobs: Optional[List[List[Any]]] = None, + ) -> ChoicesDecision: + """Select the option with the highest token length normalized prompt logprob.""" + best_choice = choices[np.argmax(normalized_prompt_logprobs)] + meta_info = { + "normalized_prompt_logprobs": normalized_prompt_logprobs, + "input_token_logprobs": input_token_logprobs, + "output_token_logprobs": output_token_logprobs, + } + return ChoicesDecision(decision=best_choice, meta_info=meta_info) + + +token_length_normalized = TokenLengthNormalized() + + +class GreedyTokenSelection(ChoicesSamplingMethod): + + def __call__( + self, + *, + choices: List[str], + normalized_prompt_logprobs: List[float], + input_token_logprobs: List[List[Any]], + output_token_logprobs: List[List[Any]], + unconditional_token_logprobs: Optional[List[List[Any]]] = None, + ) -> ChoicesDecision: + """Select the option based on greedy logprob selection. For overlapping options + where one option is a subset of a longer option, extend the shorter option using + its average logprob for comparison against the longer option.""" + + num_options = len(choices) + max_tokens = max(len(option) for option in input_token_logprobs) + logprob_matrix = self._build_logprob_matrix( + input_token_logprobs, max_tokens, num_options + ) + remaining = self._greedy_selection(logprob_matrix, num_options, max_tokens) + + best_choice = choices[remaining[0]] + meta_info = { + "normalized_prompt_logprobs": normalized_prompt_logprobs, + "input_token_logprobs": input_token_logprobs, + "output_token_logprobs": output_token_logprobs, + "greedy_logprob_matrix": logprob_matrix.tolist(), + } + return ChoicesDecision(decision=best_choice, meta_info=meta_info) + + def _build_logprob_matrix(self, input_token_logprobs, max_tokens, num_options): + logprob_matrix = np.zeros((num_options, max_tokens)) + for i, option in enumerate(input_token_logprobs): + actual_logprobs = [token[0] for token in option] + avg_logprob = np.mean(actual_logprobs) + logprob_matrix[i, : len(option)] = actual_logprobs + if len(option) < max_tokens: + logprob_matrix[i, len(option) :] = avg_logprob + return logprob_matrix + + def _greedy_selection(self, logprob_matrix, num_options, max_tokens): + remaining = np.arange(num_options) + for j in range(max_tokens): + max_logprob = np.max(logprob_matrix[remaining, j]) + remaining = remaining[logprob_matrix[remaining, j] == max_logprob] + if len(remaining) == 1: + break + return remaining + + +greedy_token_selection = GreedyTokenSelection() + + +class UnconditionalLikelihoodNormalized(ChoicesSamplingMethod): + + @property + def requires_unconditional_logprobs(self) -> bool: + return True + + def __call__( + self, + *, + choices: List[str], + normalized_prompt_logprobs: List[float], + input_token_logprobs: List[List[Any]], + output_token_logprobs: List[List[Any]], + unconditional_token_logprobs: Optional[List[List[Any]]] = None, + ) -> ChoicesDecision: + """Select the option with the highest average token logprob once normalized by + the unconditional token logprobs. + + The first unconditional token logprob is assumed to be None. If so, it is + replaced with 0 for the purposes of normalization.""" + + if unconditional_token_logprobs is None: + raise ValueError( + "Unconditional token logprobs are required for this method." + ) + + normalized_unconditional_prompt_logprobs = self._normalize_logprobs( + input_token_logprobs, unconditional_token_logprobs + ) + + best_choice = choices[np.argmax(normalized_unconditional_prompt_logprobs)] + meta_info = { + "normalized_prompt_logprobs": normalized_prompt_logprobs, + "input_token_logprobs": input_token_logprobs, + "output_token_logprobs": output_token_logprobs, + "unconditional_token_logprobs": unconditional_token_logprobs, + "normalized_unconditional_prompt_logprobs": normalized_unconditional_prompt_logprobs, + } + return ChoicesDecision(decision=best_choice, meta_info=meta_info) + + def _normalize_logprobs(self, input_token_logprobs, unconditional_token_logprobs): + normalized_unconditional_prompt_logprobs = [] + for inputs, unconditionals in zip( + input_token_logprobs, unconditional_token_logprobs + ): + inputs_logprobs = np.array([token[0] for token in inputs]) + unconditionals_logprobs = np.array([token[0] for token in unconditionals]) + unconditionals_logprobs[0] = unconditionals_logprobs[0] or 0 + normalized_unconditional_prompt_logprobs.append( + float(np.mean(inputs_logprobs - unconditionals_logprobs)) + ) + return normalized_unconditional_prompt_logprobs + + +unconditional_likelihood_normalized = UnconditionalLikelihoodNormalized() diff --git a/python/sglang/lang/compiler.py b/python/sglang/lang/compiler.py new file mode 100644 index 000000000..1284232f7 --- /dev/null +++ b/python/sglang/lang/compiler.py @@ -0,0 +1,231 @@ +import multiprocessing +from concurrent.futures import ThreadPoolExecutor +from queue import Queue +from typing import List, Union + +from sglang.global_config import global_config +from sglang.lang.interpreter import ProgramState, StreamExecutor, cache_program +from sglang.lang.ir import SglArgument, SglExpr, SglSamplingParams, SglVariable + + +def compile_func(function, backend): + tracer = function.trace(backend=backend) + compiler = CompiledFunction(tracer, function) + return compiler + + +class CompiledFunction: + def __init__(self, tracer, function): + self.function = function + + self.last_node = CompGraphNode(tracer.last_node) + self.expr_to_node = {} + self.build_graph(tracer) + self.topological_sort() + + def build_graph(self, tracer): + self.nodes = [self.last_node] + self.expr_to_node[tracer.last_node] = self.nodes[-1] + + rename_pid = {} + + visited = set([tracer.last_node]) + head = 0 + while head < len(self.nodes): + cur_node = self.nodes[head] + + # add prev node + prev_node = cur_node.expr.prev_node + if prev_node is not None: + if prev_node not in visited: + visited.add(prev_node) + self.nodes.append(CompGraphNode(prev_node)) + self.expr_to_node[prev_node] = self.nodes[-1] + cur_node.prev_node = self.expr_to_node[prev_node] + self.expr_to_node[prev_node].add_next_node(cur_node) + + # add source node + if isinstance(cur_node.expr, SglVariable): + if cur_node.expr.name in tracer.variables: + source = tracer.variables[cur_node.expr.name].source + else: + source = cur_node.expr.source + if source not in visited: + visited.add(source) + self.nodes.append(CompGraphNode(source)) + self.expr_to_node[source] = self.nodes[-1] + cur_node.source_node = self.expr_to_node[source] + self.expr_to_node[source].add_next_node(cur_node) + head += 1 + + # rename pid + if cur_node.expr.pid not in rename_pid: + rename_pid[cur_node.expr.pid] = len(rename_pid) + cur_node.expr.pid = rename_pid[cur_node.expr.pid] + + def topological_sort(self): + prevd = {} + cand = Queue() + for x in self.nodes: + prevd[x] = (x.prev_node is not None) + (x.source_node is not None) + if prevd[x] == 0: + cand.put(x) + new_list = [] + while cand.qsize() > 0: + head = cand.get() + new_list.append(head) + for x in head.next_nodes: + prevd[x] -= 1 + if prevd[x] == 0: + cand.put(x) + self.nodes = new_list + + def print_graph( + self, + ): + for node in self.nodes: + print(node) + + def run_internal( + self, + backend, + kwargs, + default_sampling_para, + ): + stream_executor_ids = set([x.expr.pid for x in self.nodes]) + stream_executors = {} + for x in stream_executor_ids: + arguments = kwargs if x == self.last_node.expr.pid else {} + stream_executors[x] = StreamExecutor( + backend, arguments, default_sampling_para, None, False + ) + for node in self.nodes: + se_id = node.expr.pid + expr = node.expr + if isinstance(expr, SglVariable): + # Make a copy for SglVariable + expr = SglVariable(expr.name, expr.source) + expr.source_stream_executor = stream_executors[ + node.source_node.expr.pid + ] + elif isinstance(expr, SglArgument): + # Substitute SglArgument + expr = kwargs[expr.name] + stream_executors[se_id].submit(expr) + for stream_executor in stream_executors.values(): + stream_executor.end() + return ProgramState(stream_executors[self.last_node.expr.pid]) + + def run( + self, + *, + max_new_tokens: int = 128, + stop: Union[str, List[str]] = (), + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + backend=None, + **kwargs, + ): + backend = backend or global_config.default_backend + + kwargs.update(self.function.bind_arguments) + + default_sampling_para = SglSamplingParams( + max_new_tokens=max_new_tokens, + stop=stop, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ) + + return self.run_internal(backend, kwargs, default_sampling_para) + + def run_batch( + self, + batch_kwargs, + *, + max_new_tokens: int = 128, + stop: Union[str, List[str]] = (), + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + backend=None, + num_threads: Union[str, int] = "auto", + ): + assert isinstance(batch_kwargs, (list, tuple)) + if len(batch_kwargs) == 0: + return [] + assert isinstance(batch_kwargs[0], dict) + + backend = backend or global_config.default_backend + + default_sampling_para = SglSamplingParams( + max_new_tokens=max_new_tokens, + stop=stop, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ) + + # Extract prefix by tracing and cache it + if len(batch_kwargs) > 1: + cache_program(self.function, backend) + + # Run all programs + if num_threads == "auto": + num_threads = multiprocessing.cpu_count() + num_threads = min(num_threads, len(batch_kwargs)) + + if num_threads == 1: + rets = [] + for arguments in batch_kwargs: + rets.append( + self.run_internal(backend, arguments, default_sampling_para) + ) + else: + with ThreadPoolExecutor(num_threads) as executor: + futures = [] + for arguments in batch_kwargs: + futures.append( + executor.submit( + self.run_internal, backend, arguments, default_sampling_para + ) + ) + rets = [f.result() for f in futures] + rets[-1].sync() + + return rets + + +class CompGraphNode: + def __init__( + self, expr: SglExpr, prev_node=None, next_nodes=None, source_node=None + ): + self.expr = expr + self.next_nodes = next_nodes or [] + self.prev_node = prev_node + self.source_node = source_node + + def add_next_node(self, other): + self.next_nodes.append(other) + + def __repr__(self): + re = f"stream {self.expr.pid:2d}: " + re += f"%{self.expr.node_id} = " + if self.prev_node is not None: + re += f"%{self.prev_node.expr.node_id} + " + re += repr(self.expr) + return re diff --git a/python/sglang/lang/interpreter.py b/python/sglang/lang/interpreter.py new file mode 100644 index 000000000..8b8cdf9c5 --- /dev/null +++ b/python/sglang/lang/interpreter.py @@ -0,0 +1,1060 @@ +"""The interpreter that executes SGL programs""" + +import asyncio +import contextvars +import copy +import multiprocessing +import queue +import threading +import uuid +import warnings +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from typing import Any, Callable, Dict, List, Optional + +import tqdm + +from sglang.global_config import global_config +from sglang.lang.ir import ( + SglCommitLazy, + SglConcateAndAppend, + SglConstantText, + SglExpr, + SglExprList, + SglGen, + SglImage, + SglRoleBegin, + SglRoleEnd, + SglSelect, + SglSeparateReasoning, + SglVariable, + SglVarScopeBegin, + SglVarScopeEnd, + SglVideo, +) +from sglang.utils import ( + encode_image_base64, + encode_video_base64, + get_exception_traceback, +) + + +def run_internal(state, program, func_args, func_kwargs, sync): + try: + state.ret_value = program.func(state, *func_args, **func_kwargs) + except Exception as e: + raise e + finally: + state.stream_executor.end() + + if sync: + state.stream_executor.sync() + + if global_config.verbosity >= 2: + print(state.text()) + + +def run_program( + program, + backend, + func_args, + func_kwargs, + default_sampling_para, + stream, + sync=False, + use_thread=True, +): + if hasattr(backend, "endpoint"): + backend = backend.endpoint + assert backend is not None, "Please specify a backend" + func_kwargs.update(program.bind_arguments) + stream_executor = StreamExecutor( + backend, + func_kwargs, + default_sampling_para, + chat_template=None, + stream=stream, + num_api_spec_tokens=program.num_api_spec_tokens, + use_thread=use_thread, + ) + state = ProgramState(stream_executor) + + if stream: + t = threading.Thread( + target=run_internal, args=(state, program, func_args, func_kwargs, sync) + ) + t.start() + return state + else: + run_internal(state, program, func_args, func_kwargs, sync) + return state + + +def run_program_batch( + program, + backend, + batch_arguments, + default_sampling_para, + num_threads, + progress_bar, + generator_style=False, +): + if hasattr(backend, "endpoint"): + backend = backend.endpoint + + # Pre-cache the common prefix for a batch. The prefix is extracted by tracing the program. + if global_config.enable_precache_with_tracing and len(batch_arguments) > 1: + cache_program(program, backend) + + # Run all programs + if num_threads == "auto": + num_threads = max(96, multiprocessing.cpu_count() * 16) + num_threads = min(num_threads, len(batch_arguments)) + + if generator_style: + return _run_program_batch_generator( + program, + backend, + batch_arguments, + default_sampling_para, + num_threads, + progress_bar, + ) + + # Original code path when generator_style=False + if num_threads == 1: + rets = [] + if progress_bar: + for arguments in tqdm.tqdm(batch_arguments): + rets.append( + run_program( + program, + backend, + (), + arguments, + default_sampling_para, + False, + True, + ) + ) + else: + for arguments in batch_arguments: + rets.append( + run_program( + program, + backend, + (), + arguments, + default_sampling_para, + False, + True, + ) + ) + else: + if progress_bar: + pbar = tqdm.tqdm(total=len(batch_arguments)) + + with ThreadPoolExecutor(num_threads) as executor: + futures = [] + for arguments in batch_arguments: + futures.append( + executor.submit( + run_program, + program, + backend, + (), + arguments, + default_sampling_para, + False, + True, + ) + ) + if progress_bar: + futures[-1].add_done_callback(lambda _: pbar.update()) + + rets = [f.result() for f in futures] + rets[-1].sync() + + if progress_bar: + pbar.close() + + return rets + + +def _run_program_batch_generator( + program, + backend, + batch_arguments, + default_sampling_para, + num_threads, + progress_bar, +): + """Helper function that yields results one by one using chunking to avoid overwhelming ThreadPoolExecutor.""" + if num_threads == 1: + iterator = tqdm.tqdm(batch_arguments) if progress_bar else batch_arguments + for arguments in iterator: + yield run_program( + program, + backend, + (), + arguments, + default_sampling_para, + False, + True, + ) + else: + pbar = tqdm.tqdm(total=len(batch_arguments)) if progress_bar else None + + # Process in chunks to avoid overwhelming ThreadPoolExecutor + # Otherwise, ThreadPoolExecutor.submit will block after adding certain number of tasks + # so we will never reach "yield" until all tasks are done + chunk_size = 200 + + with ThreadPoolExecutor(num_threads) as executor: + for chunk_start in range(0, len(batch_arguments), chunk_size): + chunk_end = min(chunk_start + chunk_size, len(batch_arguments)) + chunk_futures = [] + + # Submit chunk of tasks + for i in range(chunk_start, chunk_end): + future = executor.submit( + run_program, + program, + backend, + (), + batch_arguments[i], + default_sampling_para, + False, + True, + ) + if pbar: + future.add_done_callback(lambda _: pbar.update()) + chunk_futures.append(future) + + # Yield results from this chunk as they complete + for future in chunk_futures: + yield future.result() + + if pbar: + pbar.close() + + +def cache_program(program, backend): + from sglang.lang.tracer import extract_prefix_by_tracing + + prefix = extract_prefix_by_tracing(program, backend) + if prefix and len(prefix) > 64: + backend.cache_prefix(prefix) + + +class StreamExecutor: + """A stream executor that executes SGL expressions in a background thread.""" + + def __init__( + self, + backend, + arguments, + default_sampling_para, + chat_template, + stream, + num_api_spec_tokens=None, + use_thread=True, + ): + from sglang.lang.backend.base_backend import BaseBackend + + self.sid = uuid.uuid4().hex + self.backend: BaseBackend = backend + self.arguments: Dict[str, Any] = arguments + self.default_sampling_para = default_sampling_para + self.stream = stream + + self.variables = {} # Dict[name: str -> value: str] + self.variable_event = {} # Dict[name: str -> event: threading.Event] + self.meta_info = {} # Dict[name: str -> info: str] + self.is_finished = False + self.error_ = None + + # For completion + self.text_ = "" # The full text + + # For chat + self.messages_ = [] # The messages in the OpenAI API format + self.chat_template = chat_template or self.backend.get_chat_template() + self.cur_role = None + self.cur_role_begin_pos = None + + # For vision + self.images_ = [] + self.cur_images = [] + + # For fork/join + self.fork_start_text_pos = None + + # For speculative execution + self.num_api_spec_tokens = num_api_spec_tokens + self.speculated_text = "" + + # Worker thread + self.use_thread = use_thread + if self.use_thread: + self.queue = queue.Queue() + + def _run_worker_in_context(): + self._thread_worker_func() + + self.worker = threading.Thread( + target=contextvars.copy_context().run, args=(_run_worker_in_context,) + ) + self.worker.start() + + # For streaming + if stream: + self.stream_text_event = threading.Event() + self.stream_var_event = {} + else: + self.stream_text_event = None + self.stream_var_event = None + + def submit(self, expr: SglExpr): + self._init_var_event(expr) + + if self.use_thread: + self.queue.put(expr) + else: + self._execute(expr) + + def sync(self): + if self.use_thread: + self.queue.join() + + def get_var(self, name): + if name in self.variable_event: + self.variable_event[name].wait() + return self.variables[name] + + def set_var(self, name, value): + self.variables[name] = value + + def get_meta_info(self, name, timeout=None): + if name in self.variable_event: + got = self.variable_event[name].wait(timeout) + if not got: + raise TimeoutError(f"Timeout while waiting for event '{name}'") + ret = self.meta_info.get(name, None) + return ret + + def fork( + self, + size: int = 1, + position_ids_offset: Optional[List[int]] = None, + ): + if size > 1 and str(self.text_): + self.submit(SglCommitLazy()) + + self.sync() + size = int(size) + + exes = [ + StreamExecutor( + self.backend, + self.arguments, + self.default_sampling_para, + self.chat_template, + self.stream, + ) + for _ in range(size) + ] + for i in range(size): + exes[i].variables = dict(self.variables) + exes[i].text_ = str(self.text_) + exes[i].messages_ = list(self.messages_) + exes[i].cur_role = self.cur_role + exes[i].cur_role_begin_pos = self.cur_role_begin_pos + exes[i].fork_start_text_pos = len(self.text_) + exes[i].images_ = list(self.images_) + + # TODO(ying): handle API speculative execution + + return exes + + def text(self): + self.sync() + return self.text_ + + def messages(self): + self.sync() + return self.messages_ + + def error(self): + self.sync() + return self.error_ + + def end(self): + if self.use_thread: + if self.worker.is_alive(): + self.queue.put(None) + self.backend.end_program(self) + + def _thread_worker_func(self): + error = None + + while True: + expr = self.queue.get() + if expr is None: + self.queue.task_done() + break + + try: + self._execute(expr) + except Exception as e: + warnings.warn(f"Error in stream_executor: {get_exception_traceback()}") + error = e + break + self.queue.task_done() + if self.stream_text_event: + self.stream_text_event.set() + + # Clean the queue and events + if error is not None: + try: + while True: + self.queue.task_done() + self.queue.get_nowait() + except queue.Empty: + pass + for name in self.variable_event: + self.variable_event[name].set() + if self.stream_var_event: + for name in self.stream_var_event: + self.stream_var_event[name].set() + self.error_ = error + + if self.stream_text_event: + self.stream_text_event.set() + + self.is_finished = True + + def _execute(self, other): + if isinstance(other, str): + other = SglConstantText(other) + + assert isinstance(other, SglExpr), f"{other}" + + if isinstance(other, SglConstantText): + self._execute_fill(other.value) + elif isinstance(other, SglGen): + self._execute_gen(other) + elif isinstance(other, SglSelect): + self._execute_select(other) + elif isinstance(other, SglExprList): + for x in other.expr_list: + self._execute(x) + elif isinstance(other, SglRoleBegin): + self._execute_role_begin(other) + elif isinstance(other, SglRoleEnd): + self._execute_role_end(other) + elif isinstance(other, SglImage): + self._execute_image(other) + elif isinstance(other, SglVideo): + self._execute_video(other) + elif isinstance(other, SglVariable): + self._execute_variable(other) + elif isinstance(other, SglVarScopeBegin): + self._execute_var_scope_begin(other) + elif isinstance(other, SglVarScopeEnd): + self._execute_var_scope_end(other) + elif isinstance(other, SglCommitLazy): + self._execute_commit_lazy_operations(other) + elif isinstance(other, SglConcateAndAppend): + if ( + global_config.enable_parallel_encoding + and self.backend.support_concate_and_append + ): + self._execute_concatenate_and_append_kv_cache(other) + else: + self._execute_concatenate_and_append_text(other) + elif isinstance(other, SglSeparateReasoning): + self._execute_separate_reasoning(other) + else: + raise ValueError(f"Unknown type: {type(other)}") + + def _execute_fill(self, value: str, prefix=False): + value = str(value) + + if ( + self.cur_role == "assistant" + and self.num_api_spec_tokens is not None + and self.backend.is_chat_model + and not prefix + ): + self.backend.spec_fill(value) + return + + if self.speculated_text.startswith(value): + self.speculated_text = self.speculated_text[len(value) :] + else: + self.speculated_text = "" + + self.text_ += value + + def _execute_image(self, expr: SglImage): + path = expr.path + + base64_data = encode_image_base64(path) + + self.images_.append((path, base64_data)) + self.cur_images.append((path, base64_data)) + self.text_ += self.chat_template.image_token + + def _execute_video(self, expr: SglVideo): + path = expr.path + num_frames = expr.num_frames + + base64_data = encode_video_base64(path, num_frames) + + self.images_.append((path, base64_data)) + self.cur_images.append((path, base64_data)) + self.text_ += self.chat_template.image_token + + def _spec_gen(self, sampling_params): + stop = sampling_params.stop + max_new_tokens = sampling_params.max_new_tokens + meta_info = {} + + def regen(): + nonlocal meta_info + + sampling_params.max_new_tokens = max( + sampling_params.max_new_tokens, self.num_api_spec_tokens + ) + sampling_params.stop = None + self.speculated_text, meta_info = self.backend.generate( + self, sampling_params=sampling_params + ) + + def find_stop(): + if isinstance(stop, str): + return self.speculated_text.find(stop) + elif isinstance(stop, (tuple, list)): + pos = -1 + for stop_str in stop: + stop_pos = self.speculated_text.find(stop_str) + if stop_pos != -1 and (pos == -1 or stop_pos < pos): + pos = stop_pos + return pos + else: + raise Exception("Wrong type of stop in sampling parameters.") + + if stop is None: + if len(self.speculated_text) < max_new_tokens: + regen() + comp = self.speculated_text[:max_new_tokens] + self.speculated_text = self.speculated_text[max_new_tokens:] + elif isinstance(stop, (str, list, tuple)): + if self.speculated_text == "": + regen() + stop_pos = find_stop() + if stop_pos == -1: + stop_pos = min( + sampling_params.max_new_tokens, + len(self.speculated_text), + ) + comp = self.speculated_text[:stop_pos] + self.speculated_text = self.speculated_text[stop_pos:] + else: + raise ValueError("Wrong type of stop in sampling parameters.") + + return comp, meta_info + + def _execute_gen(self, expr: SglGen): + sampling_params = self._resolve_sampling_params(expr.sampling_params) + name = expr.name + if not self.stream: + if self.num_api_spec_tokens is None: + comp, meta_info = self.backend.generate( + self, + sampling_params=sampling_params, + ) + + else: + if self.backend.is_chat_model: + # Speculative execution on models with only chat interface. + # Store the calls into a temporary list. + # They will be lazily executed later. + comp, meta_info = self.backend.generate( + self, + sampling_params=sampling_params, + spec_var_name=name, + ) + return + + else: # Speculative execution on models with completion interface + comp, meta_info = self._spec_gen(sampling_params) + if isinstance(comp, list): + self.text_ += comp[0] + else: + assert isinstance(comp, str) + self.text_ += comp + + self.variables[name] = comp + self.meta_info[name] = meta_info + self.variable_event[name].set() + else: + assert ( + self.num_api_spec_tokens is None + ), "stream is not supported with api speculative execution" + generator = self.backend.generate_stream( + self, sampling_params=sampling_params + ) + + self.variables[name] = "" + self.stream_var_event[name].set() + + for comp, meta_info in generator: + self.text_ += comp + self.variables[name] += comp + self.meta_info[name] = meta_info + self.stream_var_event[name].set() + self.stream_text_event.set() + + self.variable_event[name].set() + self.stream_var_event[name].set() + + def _execute_select(self, expr: SglSelect): + choices_decision = self.backend.select( + self, expr.choices, expr.temperature, expr.choices_method + ) + if expr.name is not None: + name = expr.name + self.variables[name] = choices_decision.decision + self.meta_info[name] = choices_decision.meta_info + self.variable_event[name].set() + if self.stream_var_event: + self.stream_var_event[name].set() + self.text_ += choices_decision.decision + + def _execute_variable(self, expr: SglVariable): + src_executor = expr.source_stream_executor + value = src_executor.get_var(expr.name) + self._execute_fill(value) + + def _execute_role_begin(self, expr: SglRoleBegin): + assert self.cur_role is None, "Nested roles are not allowed." + + if len(self.messages_) == 0 and expr.role != "system": + # Insert the default system message + default_system = self.chat_template.default_system_prompt + if default_system: + self._execute_role_begin(SglRoleBegin("system")) + self._execute_fill(default_system) + self._execute_role_end(SglRoleEnd("system")) + + self.cur_role = expr.role + + prefix, _ = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_) + + self._execute_fill(prefix, prefix=True) + self.cur_role_begin_pos = len(self.text_) + + def _execute_role_end(self, expr: SglRoleEnd): + if ( + self.cur_role == "assistant" + and self.num_api_spec_tokens is not None + and self.backend.is_chat_model + ): + # Execute the stored lazy generation calls + self.backend.role_end_generate(self) + self.cur_role = None + + new_text = self.text_[self.cur_role_begin_pos :].lstrip() + + _, suffix = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_) + self._execute_fill(suffix) + + if self.cur_images: + # OpenAI vision API format + last_msg = { + "role": expr.role, + "content": [{"type": "text", "text": new_text}], + } + for image_path, image_base64_data in self.cur_images: + last_msg["content"].append( + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{image_base64_data}" + }, + } + ) + self.messages_.append(last_msg) + self.cur_images = [] + else: + # OpenAI chat API format + self.messages_.append({"role": expr.role, "content": new_text}) + + def _execute_var_scope_begin(self, expr: SglVarScopeBegin): + self.variables[expr.name] = int(len(self.text_)) + + def _execute_var_scope_end(self, expr: SglVarScopeEnd): + self.variables[expr.name] = self.text_[self.variables[expr.name] :] + self.variable_event[expr.name].set() + + def _execute_commit_lazy_operations(self, expr: SglCommitLazy): + self.backend.commit_lazy_operations(self) + + def _execute_concatenate_and_append_text(self, expr: SglConcateAndAppend): + new_text = "" + for s in expr.states: + exe = s.stream_executor + exe.sync() + new_text += exe.text_[exe.fork_start_text_pos :] + + self._execute_fill(new_text) + + def _execute_concatenate_and_append_kv_cache(self, expr: SglConcateAndAppend): + self_len = len(self.text_) + + for i, s in enumerate(expr.states): + exe = s.stream_executor + exe.submit(SglCommitLazy()) + + for i, s in enumerate(expr.states): + exe = s.stream_executor + exe.sync() + assert exe.fork_start_text_pos == self_len + self.text_ += exe.text_[exe.fork_start_text_pos :] + + src_rids = [state.stream_executor.sid for state in expr.states] + self.backend.concatenate_and_append(src_rids, self.sid) + + def _execute_separate_reasoning(self, expr: SglSeparateReasoning): + if self.stream: + # separate reasoning for stream is not supported + return + + if ( + self.cur_role == "assistant" + and self.num_api_spec_tokens is not None + and self.backend.is_chat_model + ): + # Execute the stored lazy generation calls + self.backend.role_end_generate(self) + + from sglang.srt.parser.reasoning_parser import ReasoningParser + + reasoning_parser = ReasoningParser(expr.model_type) + other = expr.expr + if not other: + return + elif isinstance(other, SglGen) or isinstance(other, SglSelect): + cur_text = self.get_var(other.name) + reasoning, normal_text = reasoning_parser.parse_non_stream(cur_text) + reasoning_name = expr.process_name_for_reasoning(other.name) + self.set_var(other.name, normal_text) + self.set_var(reasoning_name, reasoning) + # the variable is ready to be used + self.variable_event[reasoning_name].set() + self.text_ = self.text_[: self.cur_role_begin_pos] + normal_text + elif isinstance(other, SglExprList): + for x in other.expr_list: + self._execute_separate_reasoning( + SglSeparateReasoning(expr.model_type, x) + ) + + def _init_var_event(self, expr): + if isinstance( + expr, (SglGen, SglSelect, SglVarScopeBegin, SglSeparateReasoning) + ): + self.variable_event[expr.name] = threading.Event() + if self.stream: + self.stream_var_event[expr.name] = threading.Event() + elif isinstance(expr, SglExprList): + for e in expr.expr_list: + self._init_var_event(e) + + def _resolve_sampling_params(self, sampling_params): + """ + Construct sampling param based on default + override values + + The default values of sampling are populated in `default_sampling_para` via sgl.function.run(...sampling_args) + , and `sampling_params` contains the override values from sgl.gen(). + + Here we use default_sampling_para as the base and override the values if they exist in `sampling_params`. + It also extends the stop tokens based on the chat template. + """ + + # deepcopy is required because the dict has lists inside + clone = copy.deepcopy(self.default_sampling_para) + + for item in [ + "max_new_tokens", + "min_new_tokens", + "n", + "stop", + "stop_token_ids", + "temperature", + "top_p", + "top_k", + "min_p", + "frequency_penalty", + "presence_penalty", + "ignore_eos", + "return_logprob", + "logprob_start_len", + "top_logprobs_num", + "return_text_in_logprobs", + "dtype", + "regex", + "json_schema", + ]: + value = getattr(sampling_params, item, None) + if value is not None: + setattr(clone, item, value) + + if self.chat_template.stop_str: + if clone.stop == (): + clone.stop = [] + elif isinstance(clone.stop, str): + clone.stop = [clone.stop] + clone.stop += self.chat_template.stop_str + + return clone + + def __del__(self): + self.end() + + +class ProgramState: + """The state of an SGL program.""" + + def __init__(self, stream_executor: StreamExecutor): + self.stream_executor = stream_executor + + def _role_common(self, name: str, expr: Optional[SglExpr] = None): + if expr is not None: + role_expr = SglExprList([SglRoleBegin(name), expr, SglRoleEnd(name)]) + self.stream_executor.submit(role_expr) + return role_expr + else: + + @contextmanager + def role_scope(): + self.stream_executor.submit(SglRoleBegin(name)) + yield + self.stream_executor.submit(SglRoleEnd(name)) + + return role_scope() + + def system(self, expr: Optional[SglExpr] = None): + return self._role_common("system", expr) + + def user(self, expr: Optional[SglExpr] = None): + return self._role_common("user", expr) + + def assistant(self, expr: Optional[SglExpr] = None): + return self._role_common("assistant", expr) + + @contextmanager + def var_scope(self, name: str): + self.stream_executor.submit(SglVarScopeBegin(name)) + yield + self.stream_executor.submit(SglVarScopeEnd(name)) + + def fork( + self, + size: int = 1, + position_ids_offset: Optional[List[int]] = None, + ): + stream_executors = self.stream_executor.fork(size, position_ids_offset) + states = [ProgramState(x) for x in stream_executors] + state_group = ProgramStateGroup(states, self) + return state_group + + @contextmanager + def copy(self, position_ids_offset: Optional[List[int]] = None): + state_group = self.fork(1, position_ids_offset) + try: + yield state_group[0] + finally: + state_group.join() + + def text(self): + return self.stream_executor.text() + + def messages(self): + return self.stream_executor.messages() + + def sync(self): + return self.stream_executor.sync() + + def error(self): + return self.stream_executor.error() + + def text_iter(self, var_name: Optional[str] = None): + if self.stream_executor.stream: + prev = 0 + if var_name is None: + event = self.stream_executor.stream_text_event + while True: + event.wait() + event.clear() + out = str(self.stream_executor.text_[prev:]) + prev += len(out) + if out: + yield out + if self.stream_executor.is_finished: + break + else: + event = None + while not event: + if var_name in self.stream_executor.stream_var_event: + event = self.stream_executor.stream_var_event[var_name] + if self.stream_executor.is_finished: + yield "" + return + + while True: + event.wait() + event.clear() + out = str(self.stream_executor.variables[var_name][prev:]) + prev += len(out) + if out: + yield out + if self.stream_executor.variable_event[var_name].is_set(): + break + else: + if var_name is None: + yield self.text() + else: + yield self.get_var(var_name) + + async def text_async_iter( + self, var_name: Optional[str] = None, return_meta_data: bool = False + ): + loop = asyncio.get_running_loop() + + if self.stream_executor.stream: + prev = 0 + if var_name is None: + event = self.stream_executor.stream_text_event + while True: + await loop.run_in_executor(None, event.wait) + event.clear() + out = str(self.stream_executor.text_[prev:]) + prev += len(out) + if out: + yield out + if self.stream_executor.is_finished: + break + else: + event = None + while not event: + if var_name in self.stream_executor.stream_var_event: + event = self.stream_executor.stream_var_event[var_name] + if self.stream_executor.is_finished: + yield "" + return + + while True: + await loop.run_in_executor(None, event.wait) + event.clear() + out = str(self.stream_executor.variables[var_name][prev:]) + prev += len(out) + if out: + if return_meta_data: + yield out, self.stream_executor.meta_info[var_name] + else: + yield out + if self.stream_executor.variable_event[var_name].is_set(): + break + else: + if var_name is None: + yield self.text() + else: + yield self.get_var(var_name) + + def get_var(self, name): + return self.stream_executor.get_var(name) + + def set_var(self, name, value): + return self.stream_executor.set_var(name, value) + + def get_meta_info(self, name): + return self.stream_executor.get_meta_info(name) + + def __iadd__(self, other): + if other is None: + raise ValueError("Tried to append None to state.") + self.stream_executor.submit(other) + return self + + def __getitem__(self, name): + return self.get_var(name) + + def __setitem__(self, name, value): + self.set_var(name, value) + + def __contains__(self, name): + return name in self.stream_executor.variables + + def __del__(self): + self.stream_executor.end() + + def __repr__(self) -> str: + return f"ProgramState({self.text()})" + + +class ProgramStateGroup: + def __init__( + self, states: List[ProgramState], src_state: Optional[ProgramState] = None + ): + self.states = states + self.src_state = src_state + + def join(self, mode: str = "gather_variable"): + if mode == "gather_variable": + # Copy variables back + src_vars = self.src_state.stream_executor.variables + src_var_set = set(src_vars.keys()) + for child_state in self.states: + child_state.stream_executor.sync() + child_vars = child_state.stream_executor.variables + new_vars = set(child_vars.keys()) - src_var_set + + for k in new_vars: + if k in src_vars: + src_vars[k].append(child_vars[k]) + else: + src_vars[k] = [child_vars[k]] + elif mode == "concate_and_append": + # Concatenate and append KV cache + self.src_state += SglConcateAndAppend(self.states) + # Need a sync here. Otherwise, `states` can be deleted. + self.src_state.stream_executor.sync() + else: + raise ValueError(f"Invalid join mode: {mode}") + + for s in self.states: + s.stream_executor.end() + + def __getitem__(self, i: int): + return self.states[i] + + def __setitem__(self, i: int, value): + assert self.states[i] == value + + def __iadd__(self, other): + if isinstance(other, Callable): + # lambda function + for i in range(len(self.states)): + self.states[i] += other(i) + elif isinstance(other, SglExpr): + for i in range(len(self.states)): + self.states[i] += other + elif isinstance(other, (list, tuple)): + for i in range(len(self.states)): + self.states[i] += other[i] + else: + raise ValueError(f"Invalid value: {other}") + + return self diff --git a/python/sglang/lang/ir.py b/python/sglang/lang/ir.py new file mode 100644 index 000000000..531705ebe --- /dev/null +++ b/python/sglang/lang/ir.py @@ -0,0 +1,635 @@ +"""The intermediate representation.""" + +import dataclasses +import inspect +import warnings +from typing import List, Optional, Union + +from sglang.global_config import global_config +from sglang.lang.choices import ChoicesSamplingMethod + +REGEX_INT = r"[-+]?[0-9]+[ \n]*" +REGEX_FLOAT = r"[-+]?[0-9]*\.?[0-9]+[ \n]*" +REGEX_BOOL = r"(True|False)" +REGEX_STR = r"\"[\w\d\s]*\"" # bugs with regex r"\".*\"" in interegular pkg + + +@dataclasses.dataclass +class SglSamplingParams: + max_new_tokens: int = 128 + min_new_tokens: int = 0 + n: int = 1 + stop: Union[str, List[str]] = () + stop_token_ids: Optional[List[int]] = () + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 # -1 means disable + min_p: float = 0.0 + frequency_penalty: float = 0.0 + presence_penalty: float = 0.0 + ignore_eos: bool = False + return_logprob: Optional[bool] = None + logprob_start_len: Optional[int] = (None,) + top_logprobs_num: Optional[int] = (None,) + return_text_in_logprobs: Optional[bool] = (None,) + json_schema: Optional[str] = None + + # for constrained generation, not included in to_xxx_kwargs + dtype: Optional[str] = None + regex: Optional[str] = None + + def clone(self): + return SglSamplingParams( + self.max_new_tokens, + self.min_new_tokens, + self.n, + self.stop, + self.stop_token_ids, + self.temperature, + self.top_p, + self.top_k, + self.min_p, + self.frequency_penalty, + self.presence_penalty, + self.ignore_eos, + self.return_logprob, + self.logprob_start_len, + self.top_logprobs_num, + self.return_text_in_logprobs, + self.json_schema, + ) + + def to_openai_kwargs(self): + # OpenAI does not support top_k, so we drop it here + if self.regex is not None: + warnings.warn("Regular expression is not supported in the OpenAI backend.") + return { + "max_tokens": self.max_new_tokens, + "max_completion_tokens": self.max_new_tokens, + "n": self.n, + "stop": self.stop or None, + "temperature": self.temperature, + "top_p": self.top_p, + "frequency_penalty": self.frequency_penalty, + "presence_penalty": self.presence_penalty, + } + + def to_vertexai_kwargs(self): + if self.regex is not None: + warnings.warn( + "Regular expression is not supported in the VertexAI backend." + ) + return { + "candidate_count": 1, + "max_output_tokens": self.max_new_tokens, + "stop_sequences": self.stop, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k if self.top_k > 0 else None, + } + + def to_anthropic_kwargs(self): + # Anthropic does not support frequency_penalty or presence_penalty, so we drop it here + if self.regex is not None: + warnings.warn( + "Regular expression is not supported in the Anthropic backend." + ) + return { + "max_tokens": self.max_new_tokens, + "stop_sequences": ( + self.stop if isinstance(self.stop, (list, tuple)) else [self.stop] + ), + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + } + + def to_litellm_kwargs(self): + if self.regex is not None: + warnings.warn("Regular expression is not supported in the LiteLLM backend.") + return { + "max_tokens": self.max_new_tokens, + "stop": self.stop or None, + "temperature": self.temperature, + "top_p": self.top_p, + "frequency_penalty": self.frequency_penalty, + "presence_penalty": self.presence_penalty, + } + + def to_srt_kwargs(self): + return { + "max_new_tokens": self.max_new_tokens, + "min_new_tokens": self.min_new_tokens, + "n": self.n, + "stop": self.stop, + "stop_token_ids": self.stop_token_ids, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "frequency_penalty": self.frequency_penalty, + "presence_penalty": self.presence_penalty, + "ignore_eos": self.ignore_eos, + "regex": self.regex, + "json_schema": self.json_schema, + } + + +class SglFunction: + def __init__(self, func, num_api_spec_tokens=None, bind_arguments=None): + self.func = func + self.num_api_spec_tokens = num_api_spec_tokens + self.bind_arguments = bind_arguments or {} + self.pin_prefix_rid = None + + # Parse arguments + argspec = inspect.getfullargspec(func) + assert argspec.args[0] == "s", 'The first argument must be "s"' + self.arg_names = argspec.args[1:] + self.arg_defaults = argspec.defaults if argspec.defaults is not None else [] + + def bind(self, **kwargs): + assert all(key in self.arg_names for key in kwargs) + + new_bind_dict = {**self.bind_arguments, **kwargs} + return SglFunction(self.func, bind_arguments=new_bind_dict) + + def run( + self, + *args, + max_new_tokens: int = 128, + n: int = 1, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + ignore_eos: bool = False, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, + stream: bool = False, + backend=None, + use_thread: bool = True, + **kwargs, + ): + from sglang.lang.interpreter import run_program + + # avoid using [] as the default arg: https://nikos7am.com/posts/mutable-default-arguments/ + if stop is None: + stop = [] + if stop_token_ids is None: + stop_token_ids = [] + + default_sampling_para = SglSamplingParams( + max_new_tokens=max_new_tokens, + n=n, + stop=stop, + stop_token_ids=stop_token_ids, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ignore_eos=ignore_eos, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + return_text_in_logprobs=return_text_in_logprobs, + ) + backend = backend or global_config.default_backend + return run_program( + self, + backend, + args, + kwargs, + default_sampling_para, + stream, + use_thread=use_thread, + ) + + def run_batch( + self, + batch_kwargs, + *, + max_new_tokens: int = 128, + n: int = 1, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + ignore_eos: bool = False, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, + backend=None, + num_threads: Union[str, int] = "auto", + progress_bar: bool = False, + generator_style: bool = False, + ): + from sglang.lang.interpreter import run_program_batch + + if stop is None: + stop = [] + if stop_token_ids is None: + stop_token_ids = [] + + assert isinstance(batch_kwargs, (list, tuple)) + if len(batch_kwargs) == 0: + return [] + if not isinstance(batch_kwargs[0], dict): + num_programs = len(batch_kwargs) + # change the list of argument values to dict of arg_name -> arg_value + batch_kwargs = [ + {self.arg_names[i]: v for i, v in enumerate(arg_values)} + for arg_values in batch_kwargs + if isinstance(arg_values, (list, tuple)) + and len(self.arg_names) - len(self.arg_defaults) + <= len(arg_values) + <= len(self.arg_names) + ] + # Ensure to raise an exception if the number of arguments mismatch + if len(batch_kwargs) != num_programs: + raise Exception("Given arguments mismatch the SGL function signature") + + default_sampling_para = SglSamplingParams( + max_new_tokens=max_new_tokens, + n=n, + stop=stop, + stop_token_ids=stop_token_ids, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ignore_eos=ignore_eos, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + return_text_in_logprobs=return_text_in_logprobs, + ) + backend = backend or global_config.default_backend + return run_program_batch( + self, + backend, + batch_kwargs, + default_sampling_para, + num_threads, + progress_bar, + generator_style=generator_style, + ) + + def trace(self, *, backend=None, **kwargs): + from sglang.lang.tracer import trace_program + + backend = backend or global_config.default_backend + return trace_program(self, kwargs, backend) + + def cache(self, backend=None): + from sglang.lang.interpreter import cache_program + + backend = backend or global_config.default_backend + return cache_program(self, backend) + + def compile(self, *, backend=None): + from sglang.lang.compiler import compile_func + + return compile_func(self, backend) + + def __call__(self, *args, **kwargs): + from sglang.lang.tracer import TracingScope + + tracing_scope = TracingScope.get_current_scope() + if tracing_scope is None: + return self.run(*args, **kwargs) + else: + kwargs["backend"] = tracing_scope.tracer_state.backend + return self.trace(*args, **kwargs) + + +class SglExpr: + node_ct = 0 + + def __init__(self): + self.node_id = SglExpr.node_ct + self.prev_node = None + self.pid = None + SglExpr.node_ct += 1 + + def __add__(self, other): + if isinstance(other, str): + other = SglConstantText(other) + assert isinstance(other, SglExpr) + + return self.concatenate_ir(self, other) + + def __radd__(self, other): + if isinstance(other, str): + other = SglConstantText(other) + assert isinstance(other, SglExpr), f"{other}" + + return self.concatenate_ir(other, self) + + def concatenate_ir(self, a, b): + if isinstance(a, SglExprList): + if isinstance(b, SglExprList): + return SglExprList(a.expr_list + b.expr_list) + else: + return SglExprList(a.expr_list + [b]) + elif isinstance(b, SglExprList): + return SglExprList([a] + b.expr_list) + + return SglExprList([a, b]) + + def print_graph_dfs(self): + ret = [""] + visited = set() + + def dfs_print(x): + if x is None or x in visited: + return + visited.add(x) + + # Print dependency + if x.prev_node is not None: + dfs_print(x.prev_node) + + if isinstance(x, SglExprList): + for y in x.expr_list: + dfs_print(y) + # elif isinstance(x, SglRole): + # dfs_print(x.expr) + elif isinstance(x, SglVariable): + dfs_print(x.source) + + # Print the node itself + if isinstance(x, (SglFork, SglGetForkItem)): + ret[0] += f"%{x.node_id} = {x}\n" + else: + if x.prev_node is not None: + ret[0] += ( + f"%{x.node_id} = %{x.prev_node.node_id} + " + str(x) + "\n" + ) + else: + ret[0] += f"%{x.node_id} = " + str(x) + "\n" + + dfs_print(self) + return ret[0] + + +class SglExprList(SglExpr): + def __init__(self, expr_list: List[SglExpr]): + super().__init__() + self.expr_list = expr_list + + def __repr__(self): + return f"ExprList({self.expr_list})" + + +class SglArgument(SglExpr): + def __init__(self, name: str, value: str): + super().__init__() + self.name = name + self.value = value + + def __repr__(self): + return f"Argument(name={self.name}, value={repr(self.value)})" + + def __len__(self): + return len(self.value) + + def __getitem__(self, i): + return self.value[i] + + def __int__(self): + return self.value + + def __bool__(self): + return self.value + + def __format__(self, *args): + raise TypeError( + "Cannot put argument inside a f-string. " + "This is not compatible with the tracer. " + ) + + +class SglImage(SglExpr): + def __init__(self, path: str): + self.path = path + + def __repr__(self) -> str: + return f"SglImage({self.path})" + + +class SglVideo(SglExpr): + def __init__(self, path: str, num_frames: int): + self.path = path + self.num_frames = num_frames + + def __repr__(self) -> str: + return f"SglVideo({self.path}, {self.num_frames})" + + +class SglGen(SglExpr): + def __init__( + self, + name: Optional[str] = None, + max_new_tokens: Optional[int] = None, + min_new_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, + dtype: Optional[type] = None, + regex: Optional[str] = None, + json_schema: Optional[str] = None, + ): + """Call the model to generate. See the meaning of the arguments in docs/backend/sampling_params.md""" + super().__init__() + self.name = name + self.sampling_params = SglSamplingParams( + max_new_tokens=max_new_tokens, + min_new_tokens=min_new_tokens, + n=n, + stop=stop, + stop_token_ids=stop_token_ids, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ignore_eos=ignore_eos, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + return_text_in_logprobs=return_text_in_logprobs, + dtype=dtype, + regex=regex, + json_schema=json_schema, + ) + + def __repr__(self): + return f"Gen('{self.name}')" + + +class SglConstantText(SglExpr): + def __init__(self, value: str): + super().__init__() + self.value = value + + def __repr__(self): + return f"Constant({repr(self.value)})" + + +class SglRoleBegin(SglExpr): + def __init__(self, role: str): + super().__init__() + self.role = role + + def __repr__(self): + return f"RoleBegin({self.role})" + + +class SglRoleEnd(SglExpr): + def __init__(self, role: str): + super().__init__() + self.role = role + + def __repr__(self): + return f"RoleEnd({self.role})" + + +class SglSelect(SglExpr): + + def __init__( + self, + name: str, + choices: List[str], + temperature: float, + choices_method: ChoicesSamplingMethod, + ): + super().__init__() + self.name = name + self.choices = choices + self.temperature = temperature + self.choices_method = choices_method + + def __repr__(self): + return f"Select({self.name}, choices={self.choices}, choices_method={self.choices_method})" + + +class SglFork(SglExpr): + def __init__(self, number: int, position_ids_offset=None): + super().__init__() + self.number = number + self.position_ids_offset = position_ids_offset + + def __repr__(self): + return ( + f"Fork(%{self.prev_node.node_id}, number={self.number}, " + f"position_ids_offset={self.position_ids_offset})" + ) + + +class SglGetForkItem(SglExpr): + def __init__(self, index: int): + super().__init__() + self.index = index + + def __repr__(self): + return f"GetForkItem(%{self.prev_node.node_id}, index={self.index})" + + +class SglVariable(SglExpr): + def __init__(self, name: str, source): + super().__init__() + self.name = name + self.source = source + + def __repr__(self): + return f"Variable('{self.name}', source=%{self.source.node_id})" + + +class SglVarScopeBegin(SglExpr): + def __init__(self, name: str): + super().__init__() + self.name = name + + def __repr__(self): + return f"VarScopeBegin('{self.name}')" + + +class SglVarScopeEnd(SglExpr): + def __init__(self, name: str): + super().__init__() + self.name = name + + def __repr__(self): + return f"VarScopeEnd('{self.name}')" + + +class SglConcateAndAppend(SglExpr): + def __init__(self, states): + super().__init__() + self.states = states + + def __repr__(self): + return f"ConcatenateAndAppend('{self.states}')" + + +class SglCommitLazy(SglExpr): + def __init__(self): + super().__init__() + + def __repr__(self): + return "CommitLazy()" + + +class SglSeparateReasoning(SglExpr): + def __init__(self, model_type: str, expr: SglExpr): + super().__init__() + self.model_type = model_type + + self.expr = expr + self.name = None + self._process_expr(expr) + + def process_name_for_reasoning(self, name): + if not name: + raise ValueError("name must be provided") + return f"{name}_reasoning_content" + + def _process_expr(self, expr): + if isinstance(expr, SglGen): + self.name = self.process_name_for_reasoning(expr.name) + elif isinstance(expr, SglSelect): + self.name = self.process_name_for_reasoning(expr.name) + elif isinstance(expr, SglExprList): + for x in expr.expr_list: + self._process_expr(x) + + def __repr__(self): + return f"SeparateReasoning(model_type={self.model_type}, name={self.name})" diff --git a/python/sglang/lang/tracer.py b/python/sglang/lang/tracer.py new file mode 100644 index 000000000..0a2a744f9 --- /dev/null +++ b/python/sglang/lang/tracer.py @@ -0,0 +1,279 @@ +"""Tracing a program.""" + +import uuid +from typing import Any, Dict, List, Optional + +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.interpreter import ProgramState, ProgramStateGroup +from sglang.lang.ir import ( + SglArgument, + SglConstantText, + SglExpr, + SglExprList, + SglFork, + SglGen, + SglGetForkItem, + SglRoleBegin, + SglRoleEnd, + SglSelect, + SglVariable, + SglVarScopeBegin, + SglVarScopeEnd, +) + + +class StopTracing(Exception): + pass + + +def extract_prefix_by_tracing(program, backend): + # Create dummy arguments + dummy_arguments = {name: SglArgument(name, None) for name in program.arg_names} + arguments = dummy_arguments + arguments.update(program.bind_arguments) + + # Trace + tracer = TracerProgramState(backend, arguments, only_trace_prefix=True) + try: + with TracingScope(tracer): + tracer.ret_value = program.func(tracer, **arguments) + except (StopTracing, TypeError, AttributeError): + # Some exceptions may not be caught + pass + + # Run and cache prefix + prefix = "" + for expr in tracer.flatten_nodes(): + if isinstance(expr, SglConstantText): + prefix += expr.value + else: + break + return prefix + + +def trace_program(program, arguments, backend): + # Create dummy backend + if backend is None: + backend = BaseBackend() + + # Create dummy arguments + dummy_arguments = { + name: SglArgument(name, None) + for name in program.arg_names + if name not in arguments + } + arguments.update(dummy_arguments) + arguments.update(program.bind_arguments) + + # Trace + tracer = TracerProgramState(backend, arguments, only_trace_prefix=False) + with TracingScope(tracer): + tracer.ret_value = program.func(tracer, **arguments) + return tracer + + +class TracerProgramState(ProgramState): + def __init__(self, backend, arguments, only_trace_prefix): + self.pid = uuid.uuid4().hex + self.backend = backend + self.arguments: Dict[str, Any] = arguments + self.only_trace_prefix = only_trace_prefix + + if hasattr(backend, "endpoint"): + self.backend = backend.endpoint + + self.nodes = [] + self.last_node = None + self.variables = {} + self.ret_value = None + + # For completion + + # For chat + self.messages_ = [] + self.cur_role = None + self.chat_template = self.backend.get_chat_template() + + # For multi states + self.child_states = [] + + cur_scope = TracingScope.get_current_scope() + if cur_scope is not None: + cur_scope.add_child_state(self) + + ################################## + ########### Public API ########### + ################################## + + def fork(self, size: int = 1, position_ids_offset: Optional[List[int]] = None): + assert size >= 1 + + if self.only_trace_prefix: + raise StopTracing() + + fork_node = SglFork(size) + fork_node.prev_node = self.last_node + + states = [ + TracerProgramState(self.backend, self.arguments, self.only_trace_prefix) + for _ in range(size) + ] + + for i in range(size): + node = SglGetForkItem(i) + node.prev_node = fork_node + states[i].last_node = node + states[i].variables = dict(self.variables) + states[i].messages_ = list(self.messages_) + states[i].cur_role = self.cur_role + states[i].chat_template = self.chat_template + + state_group = ProgramStateGroup(states, self) + + return state_group + + ################################## + ########## Internal API ########## + ################################## + + def _append_node(self, other: SglExpr): + self.nodes.append(other) + other.prev_node = self.last_node + self.last_node = other + + def _execute(self, other: SglExpr): + if isinstance(other, str): + other = SglConstantText(other) + + other.pid = self.pid + + if isinstance(other, SglConstantText): + self._execute_fill(other) + elif isinstance(other, SglGen): + self._execute_gen(other) + elif isinstance(other, SglSelect): + self._execute_select(other) + elif isinstance(other, SglExprList): + for x in other.expr_list: + self._execute(x) + elif isinstance(other, SglRoleBegin): + self._execute_role_begin(other) + elif isinstance(other, SglRoleEnd): + self._execute_role_end(other) + elif isinstance(other, SglVarScopeBegin): + self._execute_var_scope_begin(other) + elif isinstance(other, SglVarScopeEnd): + self._execute_var_scope_end(other) + else: + if self.only_trace_prefix: + raise StopTracing() + else: + self._append_node(other) + + return self + + def __iadd__(self, other): + self._execute(other) + return self + + def _execute_fill(self, expr: SglConstantText): + if isinstance(expr, str): + expr = SglConstantText(expr) + self._append_node(expr) + + def _execute_gen(self, expr: SglGen): + name = expr.name if expr.name is not None else "gen_" + str(len(self.variables)) + new_node = SglVariable(name, source=expr) + self.variables[name] = new_node + self._append_node(expr) + + def _execute_select(self, expr: SglSelect): + name = ( + expr.name if expr.name is not None else "select_" + str(len(self.variables)) + ) + new_node = SglVariable(name, source=expr) + self.variables[name] = new_node + self._append_node(expr) + + def _execute_role_begin(self, expr: SglRoleBegin): + assert self.cur_role is None, "Nested roles are not allowed." + + if len(self.messages_) == 0 and expr.role != "system": + # Insert default system message + default_system = self.chat_template.default_system_prompt + if default_system: + self._execute_role_begin(SglRoleBegin("system")) + self._execute_fill(default_system) + self._execute_role_end(SglRoleEnd("system")) + + self.cur_role = expr.role + + prefix, suffix = self.chat_template.get_prefix_and_suffix( + expr.role, self.messages_ + ) + + self._execute_fill(prefix) + + def _execute_role_end(self, expr: SglRoleEnd): + prefix, suffix = self.chat_template.get_prefix_and_suffix( + expr.role, self.messages_ + ) + + self._execute_fill(suffix) + + self.messages_.append({"role": expr.role, "content": ""}) + + self.cur_role = None + + def _execute_var_scope_end(self, expr: SglVarScopeEnd): + new_node = SglVariable(expr.name, source=self.last_node) + self.variables[expr.name] = new_node + + def get_var(self, name): + ret = self.arguments.get(name, None) + if ret is not None: + return ret + + v = self.variables[name] + return SglVariable(v.name, v.source) + + def flatten_nodes(self): + def traverse(cur): + if isinstance(cur, SglExprList): + for child in cur.expr_list: + traverse(child) + else: + ret.append(cur) + + ret = [] + for x in self.nodes: + traverse(x) + return ret + + def __del__(self): + pass + + +class TracingScope: + cur_scope = None + + def __init__(self, tracer_state: TracerProgramState): + self.tracer_state = tracer_state + self.last_scope = TracingScope.cur_scope + + def __enter__(self): + TracingScope.cur_scope = self + return self + + def __exit__(self, exc_type, exc_value, traceback): + TracingScope.cur_scope = self.last_scope + + @staticmethod + def get_current_scope(): + return TracingScope.cur_scope + + def add_child_state(self, state: TracerProgramState): + cur_scope = self + while cur_scope is not None: + cur_scope.tracer_state.child_states.append(state) + cur_scope = cur_scope.last_scope diff --git a/python/sglang/launch_server.py b/python/sglang/launch_server.py new file mode 100644 index 000000000..caae7b0f6 --- /dev/null +++ b/python/sglang/launch_server.py @@ -0,0 +1,16 @@ +"""Launch the inference server.""" + +import os +import sys + +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.server_args import prepare_server_args +from sglang.srt.utils import kill_process_tree + +if __name__ == "__main__": + server_args = prepare_server_args(sys.argv[1:]) + + try: + launch_server(server_args) + finally: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/python/sglang/profiler.py b/python/sglang/profiler.py new file mode 100644 index 000000000..d872ca320 --- /dev/null +++ b/python/sglang/profiler.py @@ -0,0 +1,166 @@ +""" +Run live profiling. + +Usage: +python3 -m sglang.profiler +""" + +import argparse +import json +import os +import time +from argparse import ArgumentParser +from pathlib import Path +from typing import List, Optional + +import requests + +PARENT_FOLDER = "/tmp/sglang-profile" + + +def _run_profile( + url: Optional[str], + num_steps: int, + activities: List[str], + output_dir: Optional[str] = None, + profile_name: Optional[str] = None, + profile_by_stage: bool = False, +) -> str: + if output_dir is None: + output_dir = PARENT_FOLDER + + output_dir = os.path.normpath(output_dir) + output_dir = os.path.abspath(output_dir) + output_dir = Path(output_dir) + + # Add "profile_name/timestamp" to the path. + if profile_name: + output_dir = output_dir / profile_name + output_dir = output_dir / str(time.time()) + output_dir.mkdir(exist_ok=True, parents=True) + + print(f"Dump profiling traces to {output_dir}") + print( + f"Waiting for {num_steps} steps and the trace to be flushed.... ({profile_by_stage=})" + ) + + # Dump server args. + file_path = Path(output_dir) / "server_args.json" + if not file_path.exists(): + response = requests.get(url + "/get_server_info") + response.raise_for_status() + server_args_data = response.json() + with open(file_path, "w") as file: + file.write(json.dumps(server_args_data)) + + # Start profiler. The API replies when all steps are processed + # and files are generated. + json_data = { + "output_dir": str(output_dir), + "num_steps": str(num_steps), + "activities": activities, + "profile_by_stage": profile_by_stage, + } + + response = requests.post(url=url + "/start_profile", json=json_data) + response.raise_for_status() + + trace_link = str(output_dir) + return trace_link + + +def run_profile( + url: Optional[str], + num_steps: int, + activities: List[str], + output_dir: Optional[str] = None, + profile_name: Optional[str] = None, + profile_by_stage: bool = False, +): + # step based profile will self terminate on num_steps constraints + link = _run_profile( + url, num_steps, activities, output_dir, profile_name, profile_by_stage + ) + return link + + +if __name__ == "__main__": + parser = ArgumentParser(description="Benchmark the online serving throughput.") + parser.add_argument( + "--url", + type=str, + default="http://localhost:30000", + help="Server or API base url if not using http host and port.", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Profile directory to dump profile traces.", + ) + parser.add_argument( + "--profile-name", + type=str, + default=None, + help="The name of this profile run.", + ) + parser.add_argument( + "--num-steps", + type=int, + default=5, + help="The number of forward steps to profile.", + ) + parser.add_argument( + "--profile-by-stage", + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help="The number of forward steps to profile.", + ) + parser.add_argument( + "--cpu", + action=argparse.BooleanOptionalAction, + type=bool, + default=True, + help="Whether to profile CPU activity", + ) + parser.add_argument( + "--gpu", + action=argparse.BooleanOptionalAction, + type=bool, + default=True, + help="Whether to profile GPU activity", + ) + parser.add_argument( + "--mem", + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help="Whether to memory usage (https://pytorch.org/memory_viz)", + ) + parser.add_argument( + "--rpd", + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help="Whether to use rpd profiler (https://github.com/ROCm/rocmProfileData)", + ) + + args = parser.parse_args() + activities = [] + if args.cpu: + activities.append("CPU") + if args.gpu: + activities.append("GPU") + if args.mem: + activities.append("MEM") + if args.rpd: + activities.append("RPD") + run_profile( + args.url, + args.num_steps, + activities, + args.output_dir, + args.profile_name, + args.profile_by_stage, + ) diff --git a/python/sglang/srt/_custom_ops.py b/python/sglang/srt/_custom_ops.py new file mode 100644 index 000000000..5ed175312 --- /dev/null +++ b/python/sglang/srt/_custom_ops.py @@ -0,0 +1,177 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/_custom_ops.py +import logging +from typing import List, Optional, Tuple + +import torch + +from sglang.srt.utils import get_bool_env_var, is_hip, is_hpu, is_npu + +logger = logging.getLogger(__name__) +use_vllm_custom_allreduce = get_bool_env_var( + "USE_VLLM_CUSTOM_ALLREDUCE", default="false" +) + +if not is_hpu(): + # ROCm does not use vllm custom allreduce + if use_vllm_custom_allreduce and not is_hip(): + try: + import vllm._C + except ImportError as e: + logger.warning("Failed to import from vllm._C with %r", e) + else: + try: + import sgl_kernel + except ImportError as e: + logger.warning("Failed to import from custom_ar with %r", e) + + +if not is_hip() and not is_npu(): + if use_vllm_custom_allreduce: + custom_op = torch.ops._C_custom_ar + else: + custom_op = sgl_kernel.allreduce + + # custom allreduce + def init_custom_ar( + ipc_tensors: List[torch.Tensor], + rank_data: torch.Tensor, + rank: int, + full_nvlink: bool, + ) -> int: + return custom_op.init_custom_ar(ipc_tensors, rank_data, rank, full_nvlink) + + def all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, + ) -> None: + custom_op.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) + + def dispose(fa: int) -> None: + custom_op.dispose(fa) + + def meta_size() -> int: + return custom_op.meta_size() + + def register_buffer(fa: int, ipc_tensors: List[int]) -> None: + return custom_op.register_buffer(fa, ipc_tensors) + + def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]: + return custom_op.get_graph_buffer_ipc_meta(fa) + + def register_graph_buffers( + fa: int, handles: List[List[int]], offsets: List[List[int]] + ) -> None: + custom_op.register_graph_buffers(fa, handles, offsets) + +else: + # ROCM custom allreduce + + def init_custom_ar( + meta: torch.Tensor, + rank_data: torch.Tensor, + handles: List[str], + offsets: List[int], + rank: int, + full_nvlink: bool, + ) -> int: + return sgl_kernel.allreduce.init_custom_ar( + meta, rank_data, handles, offsets, rank, full_nvlink + ) + + def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None: + sgl_kernel.allreduce.all_reduce_reg(fa, inp, out) + + def all_reduce_unreg( + fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor + ) -> None: + sgl_kernel.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out) + + def dispose(fa: int) -> None: + sgl_kernel.allreduce.dispose(fa) + + def meta_size() -> int: + return sgl_kernel.allreduce.meta_size() + + def register_buffer( + fa: int, t: torch.Tensor, handles: List[str], offsets: List[int] + ) -> None: + return sgl_kernel.allreduce.register_buffer(fa, t, handles, offsets) + + def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]: + return sgl_kernel.allreduce.get_graph_buffer_ipc_meta(fa) + + def register_graph_buffers( + fa: int, handles: List[str], offsets: List[List[int]] + ) -> None: + sgl_kernel.allreduce.register_graph_buffers(fa, handles, offsets) + + def allocate_meta_buffer(size: int) -> torch.Tensor: + return sgl_kernel.allreduce.allocate_meta_buffer(size) + + def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor: + return sgl_kernel.allreduce.get_meta_buffer_ipc_handle(inp) + + # ROCM custom quick allreduce + + def init_custom_qr( + rank: int, world_size: int, qr_max_size: Optional[int] = None + ) -> int: + return sgl_kernel.allreduce.init_custom_qr(world_size, rank, qr_max_size) + + def qr_get_handle(fa: int) -> torch.Tensor: + return sgl_kernel.allreduce.qr_get_handle(fa) + + def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None: + sgl_kernel.allreduce.qr_open_handles(fa, handles) + + def qr_all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + quant_level: int, + cast_bf2half: bool, + ) -> None: + sgl_kernel.allreduce.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half) + + def qr_destroy(fa: int) -> None: + sgl_kernel.allreduce.qr_destroy(fa) + + def qr_max_size() -> int: + return sgl_kernel.allreduce.qr_max_size() + + +def mscclpp_generate_unique_id() -> bytes: + return sgl_kernel.allreduce.mscclpp_generate_unique_id() + + +def mscclpp_init_context( + unique_id: bytes, + rank: int, + world_size: int, + scratch: torch.Tensor, + put_buffer: torch.Tensor, + nranks_per_node: int, + rank_to_node: List[int], + rank_to_ib: List[int], + context_selection: int, +) -> int: + return sgl_kernel.allreduce.mscclpp_init_context( + unique_id, + rank, + world_size, + scratch, + put_buffer, + nranks_per_node, + rank_to_node, + rank_to_ib, + context_selection, + ) + + +def mscclpp_allreduce( + context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int +) -> None: + return sgl_kernel.allreduce.mscclpp_allreduce(context, inp, out, nthreads, nblocks) diff --git a/python/sglang/srt/aio_rwlock.py b/python/sglang/srt/aio_rwlock.py new file mode 100644 index 000000000..deda1fe79 --- /dev/null +++ b/python/sglang/srt/aio_rwlock.py @@ -0,0 +1,100 @@ +import asyncio + + +class RWLock: + def __init__(self): + # Protects internal state + self._lock = asyncio.Lock() + + # Condition variable used to wait for state changes + self._cond = asyncio.Condition(self._lock) + + # Number of readers currently holding the lock + self._readers = 0 + + # Whether a writer is currently holding the lock + self._writer_active = False + + # How many writers are queued waiting for a turn + self._waiting_writers = 0 + + @property + def reader_lock(self): + """ + A context manager for acquiring a shared (reader) lock. + + Example: + async with rwlock.reader_lock: + # read-only access + """ + return _ReaderLock(self) + + @property + def writer_lock(self): + """ + A context manager for acquiring an exclusive (writer) lock. + + Example: + async with rwlock.writer_lock: + # exclusive access + """ + return _WriterLock(self) + + async def acquire_reader(self): + async with self._lock: + # Wait until there is no active writer or waiting writer + # to ensure fairness. + while self._writer_active or self._waiting_writers > 0: + await self._cond.wait() + self._readers += 1 + + async def release_reader(self): + async with self._lock: + self._readers -= 1 + # If this was the last reader, wake up anyone waiting + # (potentially a writer or new readers). + if self._readers == 0: + self._cond.notify_all() + + async def acquire_writer(self): + async with self._lock: + # Increment the count of writers waiting + self._waiting_writers += 1 + try: + # Wait while either a writer is active or readers are present + while self._writer_active or self._readers > 0: + await self._cond.wait() + self._writer_active = True + finally: + # Decrement waiting writers only after we've acquired the writer lock + self._waiting_writers -= 1 + + async def release_writer(self): + async with self._lock: + self._writer_active = False + # Wake up anyone waiting (readers or writers) + self._cond.notify_all() + + +class _ReaderLock: + def __init__(self, rwlock: RWLock): + self._rwlock = rwlock + + async def __aenter__(self): + await self._rwlock.acquire_reader() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self._rwlock.release_reader() + + +class _WriterLock: + def __init__(self, rwlock: RWLock): + self._rwlock = rwlock + + async def __aenter__(self): + await self._rwlock.acquire_writer() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self._rwlock.release_writer() diff --git a/python/sglang/srt/bench_utils.py b/python/sglang/srt/bench_utils.py new file mode 100644 index 000000000..e9f7fcbb4 --- /dev/null +++ b/python/sglang/srt/bench_utils.py @@ -0,0 +1,137 @@ +import os +import sys +from contextlib import nullcontext + +import torch + + +# NOTE copied and modified from DeepGEMM +class suppress_stdout_stderr: + def __enter__(self): + self.outnull_file = open(os.devnull, "w") + self.errnull_file = open(os.devnull, "w") + + self.old_stdout_fileno_undup = sys.stdout.fileno() + self.old_stderr_fileno_undup = sys.stderr.fileno() + + self.old_stdout_fileno = os.dup(sys.stdout.fileno()) + self.old_stderr_fileno = os.dup(sys.stderr.fileno()) + + self.old_stdout = sys.stdout + self.old_stderr = sys.stderr + + os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup) + os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup) + + sys.stdout = self.outnull_file + sys.stderr = self.errnull_file + return self + + def __exit__(self, *_): + sys.stdout = self.old_stdout + sys.stderr = self.old_stderr + + os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup) + os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup) + + os.close(self.old_stdout_fileno) + os.close(self.old_stderr_fileno) + + self.outnull_file.close() + self.errnull_file.close() + + +# NOTE copied and modified from DeepGEMM +def bench_kineto( + fn, + kernel_names, + num_tests: int = 30, + suppress_kineto_output: bool = False, + trace_path: str = None, + flush_l2: bool = True, + with_multiple_kernels: bool = False, +): + # Conflict with Nsight Systems + using_nsys = int(os.environ.get("SGLANG_NSYS_PROFILING", 0)) + + # By default, flush L2 with an excessive 8GB memset to give the GPU some (literal) chill time without full idle + flush_l2_size = int(8e9 // 4) + + # For some auto-tuning kernels with prints + fn() + + # Profile + suppress = ( + suppress_stdout_stderr + if suppress_kineto_output and not using_nsys + else nullcontext + ) + with suppress(): + schedule = ( + torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1) + if not using_nsys + else None + ) + profiler = ( + torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule + ) + if not using_nsys + else nullcontext() + ) + with profiler: + for i in range(2): + for _ in range(num_tests): + if flush_l2: + torch.empty( + flush_l2_size, dtype=torch.int, device="cuda" + ).zero_() + fn() + + if not using_nsys: + profiler.step() + + # Return 1 if using Nsight Systems + if using_nsys: + return 1 + + # Parse the profiling table + assert isinstance(kernel_names, str) or isinstance(kernel_names, tuple) + is_tuple = isinstance(kernel_names, tuple) + prof_lines = ( + profiler.key_averages() + .table(sort_by="cuda_time_total", max_name_column_width=100) + .split("\n") + ) + kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names + assert all([isinstance(name, str) for name in kernel_names]) + if not with_multiple_kernels: + for name in kernel_names: + assert ( + sum([name in line for line in prof_lines]) == 1 + ), f"Errors of the kernel {name} in the profiling table (table: {prof_lines})" + + # Save chrome traces + if trace_path is not None: + profiler.export_chrome_trace(trace_path) + + # Return average kernel times + units = {"ms": 1e3, "us": 1e6} + kernel_times = [] + for name in kernel_names: + total_time = 0 + total_num = 0 + for line in prof_lines: + if name in line: + time_str = line.split()[-2] + num_str = line.split()[-1] + for unit, scale in units.items(): + if unit in time_str: + total_time += ( + float(time_str.replace(unit, "")) / scale * int(num_str) + ) + total_num += int(num_str) + break + kernel_times.append(total_time / total_num) + + return tuple(kernel_times) if is_tuple else kernel_times[0] diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py new file mode 100644 index 000000000..ef880c911 --- /dev/null +++ b/python/sglang/srt/configs/__init__.py @@ -0,0 +1,29 @@ +from sglang.srt.configs.chatglm import ChatGLMConfig +from sglang.srt.configs.dbrx import DbrxConfig +from sglang.srt.configs.deepseekvl2 import DeepseekVL2Config +from sglang.srt.configs.exaone import ExaoneConfig +from sglang.srt.configs.janus_pro import MultiModalityConfig +from sglang.srt.configs.kimi_vl import KimiVLConfig +from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig +from sglang.srt.configs.longcat_flash import LongcatFlashConfig +from sglang.srt.configs.qwen3_next import Qwen3NextConfig +from sglang.srt.configs.step3_vl import ( + Step3TextConfig, + Step3VisionEncoderConfig, + Step3VLConfig, +) + +__all__ = [ + "ExaoneConfig", + "ChatGLMConfig", + "DbrxConfig", + "DeepseekVL2Config", + "LongcatFlashConfig", + "MultiModalityConfig", + "KimiVLConfig", + "MoonViTConfig", + "Step3VLConfig", + "Step3TextConfig", + "Step3VisionEncoderConfig", + "Qwen3NextConfig", +] diff --git a/python/sglang/srt/configs/chatglm.py b/python/sglang/srt/configs/chatglm.py new file mode 100644 index 000000000..9370c218a --- /dev/null +++ b/python/sglang/srt/configs/chatglm.py @@ -0,0 +1,78 @@ +# Adapted from +# https://github.com/THUDM/ChatGLM2-6B +# https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/configs/chatglm.py + +# ChatGLM2 and ChatGLM3 share the same config. +# ChatGLM4 is officially supported by Huggingface +# transformers >= 4.46.0 is required +# https://huggingface.co/docs/transformers/en/model_doc/glm +from transformers import PretrainedConfig + + +class ChatGLMConfig(PretrainedConfig): + model_type = "chatglm" + attribute_map = { + "num_hidden_layers": "num_layers", + "n_head_kv": "multi_query_group_num", + } + + def __init__( + self, + num_layers=28, + padded_vocab_size=65024, + hidden_size=4096, + ffn_hidden_size=13696, + kv_channels=128, + num_attention_heads=32, + seq_length=2048, + hidden_dropout=0.0, + attention_dropout=0.0, + layernorm_epsilon=1e-5, + rmsnorm=True, + apply_residual_connection_post_layernorm=False, + post_layer_norm=True, + add_bias_linear=False, + add_qkv_bias=False, + interleaved_qkv=False, + bias_dropout_fusion=True, + multi_query_attention=False, + multi_query_group_num=1, + apply_query_key_layer_scaling=True, + attention_softmax_in_fp32=True, + fp32_residual_connection=False, + quantization_bit=0, + pre_seq_len=None, + prefix_projection=False, + **kwargs + ): + self.num_layers = num_layers + self.vocab_size = padded_vocab_size + self.padded_vocab_size = padded_vocab_size + self.hidden_size = hidden_size + self.ffn_hidden_size = ffn_hidden_size + self.kv_channels = kv_channels + self.num_attention_heads = num_attention_heads + self.seq_length = seq_length + # It is to be compatible with long lora. + self.max_position_embeddings = seq_length + self.hidden_dropout = hidden_dropout + self.attention_dropout = attention_dropout + self.layernorm_epsilon = layernorm_epsilon + self.rmsnorm = rmsnorm + self.apply_residual_connection_post_layernorm = ( + apply_residual_connection_post_layernorm + ) + self.post_layer_norm = post_layer_norm + self.add_bias_linear = add_bias_linear + self.add_qkv_bias = add_qkv_bias + self.bias_dropout_fusion = bias_dropout_fusion + self.multi_query_attention = multi_query_attention + self.multi_query_group_num = multi_query_group_num + self.apply_query_key_layer_scaling = apply_query_key_layer_scaling + self.attention_softmax_in_fp32 = attention_softmax_in_fp32 + self.fp32_residual_connection = fp32_residual_connection + self.quantization_bit = quantization_bit + self.pre_seq_len = pre_seq_len + self.prefix_projection = prefix_projection + self.interleaved_qkv = interleaved_qkv + super().__init__(**kwargs) diff --git a/python/sglang/srt/configs/dbrx.py b/python/sglang/srt/configs/dbrx.py new file mode 100644 index 000000000..75ccbde94 --- /dev/null +++ b/python/sglang/srt/configs/dbrx.py @@ -0,0 +1,279 @@ +# Adapted from +# https://huggingface.co/databricks/dbrx-base/blob/main/configuration_dbrx.py +# https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/configs/dbrx.py +"""Dbrx configuration.""" + +from typing import Any, Optional + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +DBRX_PRETRAINED_CONFIG_ARCHIVE_MAP = {} # type: ignore + + +class DbrxAttentionConfig(PretrainedConfig): + """Configuration class for Dbrx Attention. + + [`DbrxAttention`] class. It is used to instantiate attention layers + according to the specified arguments, defining the layers architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + attn_pdrop (`float`, *optional*, defaults to 0.0): + The dropout probability for the attention layers. + clip_qkv (`float`, *optional*, defaults to None): + If not `None`, clip the queries, keys, and values in the attention layer to this value. + kv_n_heads (Optional[int]): For grouped_query_attention only, allow user to specify number of kv heads. + rope_theta (float): The base frequency for rope. + """ + + def __init__( + self, + attn_pdrop: float = 0, + clip_qkv: Optional[float] = None, + kv_n_heads: int = 1, + rope_theta: float = 10000.0, + **kwargs: Any, + ): + super().__init__(**kwargs) + self.attn_pdrop = attn_pdrop + self.clip_qkv = clip_qkv + self.kv_n_heads = kv_n_heads + self.rope_theta = rope_theta + + for k in ["model_type"]: + if k in kwargs: + kwargs.pop(k) + if len(kwargs) != 0: + raise ValueError(f"Found unknown {kwargs=}") + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: str, **kwargs: Any + ) -> "PretrainedConfig": + cls._set_token_in_kwargs(kwargs) + + config_dict, kwargs = cls.get_config_dict( + pretrained_model_name_or_path, **kwargs + ) + + if config_dict.get("model_type") == "dbrx": + config_dict = config_dict["attn_config"] + + if ( + "model_type" in config_dict + and hasattr(cls, "model_type") + and config_dict["model_type"] != cls.model_type + ): + logger.warning( + "You are using a model of type %s to instantiate a model of " + "type %s. This is not supported for all configurations of " + "models and can yield errors.", + config_dict["model_type"], + cls.model_type, + ) + + return cls.from_dict(config_dict, **kwargs) + + +class DbrxFFNConfig(PretrainedConfig): + """Configuration class for Dbrx FFN. + + [`DbrxFFN`] class. It is used to instantiate feedforward layers according to + the specified arguments, defining the layers architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + ffn_act_fn (dict, optional): A dict specifying activation function for the FFN. + The dict should have a key 'name' with the value being the name of + the activation function along with any additional keyword arguments. + ffn_hidden_size (int, optional): The hidden size of the feedforward network. + moe_num_experts (int, optional): The number of experts in the mixture of experts layer. + moe_top_k (int, optional): The number of experts to use in the mixture of experts layer. + moe_jitter_eps (float, optional): The jitter epsilon for the mixture of experts layer. + moe_loss_weight (float, optional): The loss weight for the mixture of experts layer. + moe_normalize_expert_weights (float, optional): The normalization factor for the expert weights. + uniform_expert_assignment (bool, optional): Whether to use uniform expert assignment. + This should only be used for benchmarking purposes. + """ + + def __init__( + self, + ffn_act_fn: Optional[dict] = None, + ffn_hidden_size: int = 3584, + moe_num_experts: int = 4, + moe_top_k: int = 1, + moe_jitter_eps: Optional[float] = None, + moe_loss_weight: float = 0.01, + moe_normalize_expert_weights: Optional[float] = 1, + uniform_expert_assignment: bool = False, + **kwargs: Any, + ): + super().__init__() + if ffn_act_fn is None: + ffn_act_fn = {"name": "silu"} + self.ffn_act_fn = ffn_act_fn + self.ffn_hidden_size = ffn_hidden_size + self.moe_num_experts = moe_num_experts + self.moe_top_k = moe_top_k + self.moe_jitter_eps = moe_jitter_eps + self.moe_loss_weight = moe_loss_weight + self.moe_normalize_expert_weights = moe_normalize_expert_weights + self.uniform_expert_assignment = uniform_expert_assignment + + for k in ["model_type"]: + if k in kwargs: + kwargs.pop(k) + if len(kwargs) != 0: + raise ValueError(f"Found unknown {kwargs=}") + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: str, **kwargs: Any + ) -> "PretrainedConfig": + cls._set_token_in_kwargs(kwargs) + + config_dict, kwargs = cls.get_config_dict( + pretrained_model_name_or_path, **kwargs + ) + + if config_dict.get("model_type") == "dbrx": + config_dict = config_dict["ffn_config"] + + if ( + "model_type" in config_dict + and hasattr(cls, "model_type") + and config_dict["model_type"] != cls.model_type + ): + logger.warning( + "You are using a model of type %s to instantiate a model of " + "type %s. This is not supported for all " + "configurations of models and can yield errors.", + config_dict["model_type"], + cls.model_type, + ) + + return cls.from_dict(config_dict, **kwargs) + + +class DbrxConfig(PretrainedConfig): + """Configuration class for Dbrx. + + [`DbrxModel`]. It is used to instantiate a Dbrx model according to the + specified arguments, defining the model architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + d_model (`int`, *optional*, defaults to 6144): + Dimensionality of the embeddings and hidden states. + n_heads (`int`, *optional*, defaults to 48): + Number of attention heads for each attention layer in the Transformer encoder. + n_layers (`int`, *optional*, defaults to 40): + Number of hidden layers in the Transformer encoder. + max_seq_len (`int`, *optional*, defaults to 32768): + The maximum sequence length of the model. + vocab_size (`int`, *optional*, defaults to 100352): + Vocabulary size of the Dbrx model. Defines the maximum number of different tokens that can be represented by + the `inputs_ids` passed when calling [`DbrxModel`]. + resid_pdrop (`float`, *optional*, defaults to 0.0): + The dropout probability applied to the attention output before combining with residual. + emb_pdrop (`float`, *optional*, defaults to 0.0): + The dropout probability for the embedding layer. + attn_config (`dict`, *optional*): + A dictionary used to configure the model's attention module. + ffn_config (`dict`, *optional*): + A dictionary used to configure the model's FFN module. + use_cache (`bool`, *optional*, defaults to `False`): + Whether or not the model should return the last key/values attentions (not used by all models). + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + output_router_logits (`bool`, *optional*, defaults to `False`): + Whether or not the router logits should be returned by the model. Enabling this will also + allow the model to output the auxiliary loss. See [here]() for more details + router_aux_loss_coef (`float`, *optional*, defaults to 0.001): + The aux loss factor for the total loss. + + + Example: + ```python + >>> from transformers import DbrxConfig, DbrxModel + + >>> # Initializing a Dbrx configuration + >>> configuration = DbrxConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = DbrxModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "dbrx" + attribute_map = { + "num_attention_heads": "n_heads", + "hidden_size": "d_model", + "num_hidden_layers": "n_layers", + "max_position_embeddings": "max_seq_len", + } + + def __init__( + self, + d_model: int = 2048, + n_heads: int = 16, + n_layers: int = 24, + max_seq_len: int = 2048, + vocab_size: int = 32000, + resid_pdrop: float = 0.0, + emb_pdrop: float = 0.0, + attn_config: Optional[DbrxAttentionConfig] = None, + ffn_config: Optional[DbrxFFNConfig] = None, + use_cache: bool = True, + initializer_range: float = 0.02, + output_router_logits: bool = False, + router_aux_loss_coef: float = 0.05, + **kwargs: Any, + ): + if attn_config is None: + self.attn_config = DbrxAttentionConfig() + elif isinstance(attn_config, dict): + self.attn_config = DbrxAttentionConfig(**attn_config) + else: + self.attn_config = attn_config + + if ffn_config is None: + self.ffn_config = DbrxFFNConfig() + elif isinstance(ffn_config, dict): + self.ffn_config = DbrxFFNConfig(**ffn_config) + else: + self.ffn_config = ffn_config + + self.d_model = d_model + self.n_heads = n_heads + self.n_layers = n_layers + self.max_seq_len = max_seq_len + self.vocab_size = vocab_size + self.resid_pdrop = resid_pdrop + self.emb_pdrop = emb_pdrop + self.use_cache = use_cache + self.initializer_range = initializer_range + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + + tie_word_embeddings = kwargs.pop("tie_word_embeddings", False) + if tie_word_embeddings: + raise ValueError("tie_word_embeddings is not supported for Dbrx models.") + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/python/sglang/srt/configs/deepseekvl2.py b/python/sglang/srt/configs/deepseekvl2.py new file mode 100644 index 000000000..bcb0afe5a --- /dev/null +++ b/python/sglang/srt/configs/deepseekvl2.py @@ -0,0 +1,688 @@ +import math +import os +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import torch +from PIL import Image, ImageOps +from transformers import ( + AutoProcessor, + LlamaTokenizerFast, + PretrainedConfig, + ProcessorMixin, +) + + +def select_best_resolution(image_size, candidate_resolutions): + # used for cropping + original_width, original_height = image_size + best_fit = None + max_effective_resolution = 0 + min_wasted_resolution = float("inf") + + for width, height in candidate_resolutions: + scale = min(width / original_width, height / original_height) + downscaled_width, downscaled_height = int(original_width * scale), int( + original_height * scale + ) + effective_resolution = min( + downscaled_width * downscaled_height, original_width * original_height + ) + wasted_resolution = (width * height) - effective_resolution + + if effective_resolution > max_effective_resolution or ( + effective_resolution == max_effective_resolution + and wasted_resolution < min_wasted_resolution + ): + max_effective_resolution = effective_resolution + min_wasted_resolution = wasted_resolution + best_fit = (width, height) + + return best_fit + + +class DictOutput(object): + def items(self): + return self.__dict__.items() + + def keys(self): + return self.__dict__.keys() + + def __getitem__(self, item): + return self.__dict__[item] + + def __contains__(self, key): + return key in self.__dict__ + + def __setitem__(self, key, value): + self.__dict__[key] = value + + +@dataclass +class VLChatProcessorOutput(DictOutput): + input_ids: torch.LongTensor + target_ids: torch.LongTensor + pixel_values: ( + torch.Tensor + ) # rename from "images" to "pixel_values" for compatibility + images_seq_mask: torch.BoolTensor + images_spatial_crop: torch.LongTensor + + def __len__(self): + return len(self.input_ids) + + +class ImageTransform(object): + def __init__( + self, + mean: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), + std: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), + normalize: bool = True, + ): + self.mean = mean + self.std = std + self.normalize = normalize + + # only load torchvision.transforms when needed + try: + import torchvision.transforms as T + + # FIXME: add version check for gguf + except ImportError as err: + raise ImportError( + "Please install torchvision via `pip install torchvision` to use Deepseek-VL2." + ) from err + + transform_pipelines = [T.ToTensor()] + + if normalize: + transform_pipelines.append(T.Normalize(mean, std)) + + self.transform = T.Compose(transform_pipelines) + + def __call__(self, pil_img: Image.Image): + x = self.transform(pil_img) + return x + + +class DeepseekVLV2Processor(ProcessorMixin): + tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") + attributes = ["tokenizer"] + + def __init__( + self, + tokenizer: LlamaTokenizerFast, + candidate_resolutions: Tuple[Tuple[int, int]], + patch_size: int, + downsample_ratio: int, + image_mean: Tuple[float, float, float] = (0.5, 0.5, 0.5), + image_std: Tuple[float, float, float] = (0.5, 0.5, 0.5), + normalize: bool = True, + image_token: str = "", + pad_token: str = "<｜▁pad▁｜>", + add_special_token: bool = False, + sft_format: str = "deepseek", + mask_prompt: bool = True, + ignore_id: int = -100, + **kwargs, + ): + + self.candidate_resolutions = candidate_resolutions + self.image_size = candidate_resolutions[0][0] + self.patch_size = patch_size + self.image_mean = image_mean + self.image_std = image_std + self.normalize = normalize + self.downsample_ratio = downsample_ratio + + self.image_transform = ImageTransform( + mean=image_mean, std=image_std, normalize=normalize + ) + self.tokenizer = tokenizer + # must set this，padding side with make a difference in batch inference + self.tokenizer.padding_side = "left" + + # add the pad_token as special token to use 'tokenizer.pad_token' and 'tokenizer.pad_token_id' + if tokenizer.pad_token is None: + self.tokenizer.add_special_tokens({"pad_token": pad_token}) + + # add image token + image_token_id = self.tokenizer.vocab.get(image_token) + if image_token_id is None: + special_tokens = [image_token] + special_tokens_dict = {"additional_special_tokens": special_tokens} + self.tokenizer.add_special_tokens(special_tokens_dict) + self.image_token_id = self.tokenizer.vocab.get(image_token) + + # add five special tokens for grounding-related tasks + # <|ref|>, <|/ref|>, <|det|>, <|/det|>, <|grounding|> + special_tokens = ["<|ref|>", "<|/ref|>", "<|det|>", "<|/det|>", "<|grounding|>"] + special_tokens_dict = {"additional_special_tokens": special_tokens} + self.tokenizer.add_special_tokens(special_tokens_dict) + + # add special tokens for SFT data + special_tokens = ["<|User|>", "<|Assistant|>"] + special_tokens_dict = {"additional_special_tokens": special_tokens} + self.tokenizer.add_special_tokens(special_tokens_dict) + + self.image_token = image_token + self.pad_token = pad_token + self.add_special_token = add_special_token + self.sft_format = sft_format + self.mask_prompt = mask_prompt + self.ignore_id = ignore_id + + super().__init__( + tokenizer, + **kwargs, + ) + + def format_messages_v2(self, messages, pil_images, max_req_input_len=-1): + """play the role of format_messages_v2 and get_images_info in the last version""" + tokenized_data = [] + masked_tokenized_data = [] # labels + images_list = [] + images_seq_mask = [] + images_spatial_crop = [] + + image_index = 0 + image_token_cnt = messages.count(self.image_token) + tokenized_str, images, seq_mask, spatial_crop = self.tokenize_with_images( + messages, + pil_images[image_index : image_index + image_token_cnt], + bos=True, + eos=True, + cropping=len(pil_images) <= 2, + max_req_input_len=max_req_input_len, + ) + + image_index = image_token_cnt + tokenized_data += tokenized_str + if self.mask_prompt: + masked_tokenized_data += [self.ignore_id] * len(tokenized_str) + else: + masked_tokenized_data += tokenized_str + images_list += images + images_seq_mask += seq_mask + images_spatial_crop += spatial_crop + + assert len(tokenized_data) == len( + images_seq_mask + ), f"format_messages_v2: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + + return ( + tokenized_data, + masked_tokenized_data, + images_list, + images_seq_mask, + images_spatial_crop, + ) + + @property + def bos_id(self): + return self.tokenizer.bos_token_id + + @property + def eos_id(self): + return self.tokenizer.eos_token_id + + @property + def pad_id(self): + return self.tokenizer.pad_token_id + + def encode(self, text: str, bos: bool = True, eos: bool = False): + t = self.tokenizer.encode(text, add_special_tokens=False) + + if bos: + t = [self.bos_id] + t + if eos: + t = t + [self.eos_id] + + return t + + def decode(self, t: List[int], **kwargs) -> str: + return self.tokenizer.decode(t, **kwargs) + + def process_one( + self, + prompt: str = None, + conversations: List[Dict[str, str]] = None, + images: List[Image.Image] = None, + apply_sft_format: bool = False, + inference_mode: bool = True, + system_prompt: str = "", + max_req_input_len: int = -1, + **kwargs, + ): + """ + + Args: + prompt (str): the formatted prompt; + conversations (List[Dict]): conversations with a list of messages; + images (List[ImageType]): the list of images; + apply_sft_format (bool): if prompt is not None, then apply the SFT format to prompt; + if conversations is not None, then it will always apply the SFT format to conversations; + inference_mode (bool): if True, then remove the last eos token; + system_prompt (str): the system prompt; + **kwargs: + + Returns: + outputs (BaseProcessorOutput): the output of the processor, + - input_ids (torch.LongTensor): [N + image tokens] + - target_ids (torch.LongTensor): [N + image tokens] + - images (torch.FloatTensor): [n_images, 3, H, W] + - image_id (int): the id of the image token + - num_image_tokens (List[int]): the number of image tokens + """ + + assert ( + prompt is None or conversations is None + ), "prompt and conversations cannot be used at the same time." + + ( + tokenized_str, + masked_tokenized_str, + images_list, + images_seq_mask, + images_spatial_crop, + ) = self.format_messages_v2(conversations, images, max_req_input_len) + + assert ( + len(tokenized_str) == len(images_seq_mask) == len(masked_tokenized_str) + ), ( + f"tokenized_str's length {len(tokenized_str)}, input_ids' length {len(masked_tokenized_str)}, " + f"imags_seq_mask's length {len(images_seq_mask)}, are not equal" + ) + + input_ids = torch.LongTensor(tokenized_str) + target_ids = torch.LongTensor(masked_tokenized_str) + images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool) + + # set input_ids < 0 | input_ids == self.image_token_id as ignore_id + target_ids[(input_ids < 0) | (input_ids == self.image_token_id)] = ( + self.ignore_id + ) + input_ids[input_ids < 0] = self.pad_id + + if inference_mode: + assert input_ids[-1] == self.eos_id + input_ids = input_ids[:-1] + target_ids = target_ids[:-1] + images_seq_mask = images_seq_mask[:-1] + + if len(images_list) == 0: + images = torch.zeros((1, 3, self.image_size, self.image_size)) + images_spatial_crop = torch.zeros((1, 2), dtype=torch.long) + else: + images = torch.stack(images_list, dim=0) + images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long) + + images_spatial_crop = torch.stack( + [images_spatial_crop], dim=0 + ) # stack the tensor to make it a batch of 1 + + prepare = VLChatProcessorOutput( + input_ids=input_ids, + target_ids=target_ids, + pixel_values=images, + images_seq_mask=images_seq_mask, + images_spatial_crop=images_spatial_crop, + ) + + return prepare + + def __call__( + self, + *, + prompt: str = None, + conversations: List[Dict[str, str]] = None, + images: List[Image.Image] = None, + apply_sft_format: bool = False, + inference_mode: bool = True, + system_prompt: str = "", + max_req_input_len: int = -1, + **kwargs, + ): + prepare = self.process_one( + prompt=prompt, + conversations=conversations, + images=images, + apply_sft_format=apply_sft_format, + inference_mode=inference_mode, + system_prompt=system_prompt, + max_req_input_len=max_req_input_len, + ) + + return prepare + + def find_all_indices(self, messages, target_value): + indices = [] + for index, item in enumerate(messages): + if item == target_value: + indices.append(index) + return indices + + def tokenize_with_images( + self, + conversation: str, + images: List[Image.Image], + bos: bool = True, + eos: bool = True, + cropping: bool = True, + max_req_input_len: int = -1, + ): + """Tokenize text with tags.""" + images_list, images_seq_mask, images_spatial_crop = [], [], [] + text_splits = conversation.split(self.image_token) + tokenized_str = [] + for text_sep, image in zip(text_splits, images): + """encode text_sep""" + tokenized_sep = self.encode(text_sep, bos=False, eos=False) + tokenized_str += tokenized_sep + images_seq_mask += [False] * len(tokenized_sep) + + """select best resolution for anyres""" + if cropping: + best_width, best_height = select_best_resolution( + image.size, self.candidate_resolutions + ) + else: + best_width, best_height = self.image_size, self.image_size + # print(image.size, (best_width, best_height)) # check the select_best_resolutions func + + """process the global view""" + global_view = ImageOps.pad( + image, + (self.image_size, self.image_size), + color=tuple(int(x * 255) for x in self.image_transform.mean), + ) + images_list.append(self.image_transform(global_view)) + + """process the local views""" + local_view = ImageOps.pad( + image, + (best_width, best_height), + color=tuple(int(x * 255) for x in self.image_transform.mean), + ) + for i in range(0, best_height, self.image_size): + for j in range(0, best_width, self.image_size): + images_list.append( + self.image_transform( + local_view.crop( + (j, i, j + self.image_size, i + self.image_size) + ) + ) + ) + + """record height / width crop num""" + num_width_tiles, num_height_tiles = ( + best_width // self.image_size, + best_height // self.image_size, + ) + images_spatial_crop.append([num_width_tiles, num_height_tiles]) + + """add image tokens""" + h = w = math.ceil( + (self.image_size // self.patch_size) / self.downsample_ratio + ) + # global views tokens h * (w + 1), 1 is for line separator + tokenized_image = [self.image_token_id] * h * (w + 1) + # add a separator between global and local views + tokenized_image += [self.image_token_id] + # local views tokens, (num_height_tiles * h) * (num_width_tiles * w + 1) + tokenized_image += ( + [self.image_token_id] + * (num_height_tiles * h) + * (num_width_tiles * w + 1) + ) + + tokenized_str += tokenized_image + images_seq_mask += [True] * len(tokenized_image) + # print(width_crop_num, height_crop_num, len(tokenized_image)) # test the correctness of the number of image-related tokens + + """process the last text split""" + tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False) + # deal with video, limit with request len + if max_req_input_len > -1: + if max_req_input_len < len(tokenized_sep) + len(tokenized_str) - 1: + rest = max_req_input_len - len(tokenized_sep) - 1 - 1024 + tokenized_str = tokenized_str[:rest] + images_seq_mask = images_seq_mask[:rest] + tokenized_str += tokenized_sep + images_seq_mask += [False] * len(tokenized_sep) + + """add the bos and eos tokens""" + if bos: + tokenized_str = [self.bos_id] + tokenized_str + images_seq_mask = [False] + images_seq_mask + if eos: + tokenized_str = tokenized_str + [self.eos_id] + images_seq_mask = images_seq_mask + [False] + + assert len(tokenized_str) == len( + images_seq_mask + ), f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + + return tokenized_str, images_list, images_seq_mask, images_spatial_crop + + +class DeepseekVL2VisionEncoderConfig(PretrainedConfig): + model_type: str = "vision" + + model_name: str = "siglip_large_patch16_384" + image_size: int = 384 + patch_size: int = 16 + width: int = 1024 + layers: int = 24 + heads: int = 16 + mlp_ratio: int = 4 + global_pool: str = "map" + ignore_head: bool = True + class_token: bool = False + num_classes: int = 0 + use_checkpoint: bool = False + weight_init: str = "skip" + deterministic: bool = False + num_recomputing_layers: int = 0 + + def __init__( + self, + model_name: str = "siglip_large_patch16_384", + image_size: int = 384, + patch_size: int = 16, + width: int = 1024, + layers: int = 24, + heads: int = 16, + mlp_ratio: int = 4, + global_pool: str = "map", + ignore_head: bool = True, + class_token: bool = False, + num_classes: int = 0, + use_checkpoint: bool = False, + **kwargs, + ): + self.model_name = model_name + self.image_size = image_size + self.patch_size = patch_size + self.width = width + self.layers = layers + self.heads = heads + self.mlp_ratio = mlp_ratio + self.global_pool = global_pool + self.ignore_head = ignore_head + self.class_token = class_token + self.num_classes = num_classes + self.use_checkpoint = use_checkpoint + + super().__init__(**kwargs) + + +class DeepseekVL2MlpProjectorConfig(PretrainedConfig): + model_type = "mlp_projector" + projector_type: str = "downsample_mlp_gelu" + input_dim: int = 1152 + n_embed: int = 2048 + depth: int = 2 + mlp_ratio: int = 1 + downsample_ratio: int = 2 + token_pooling: bool = False + + def __init__( + self, + projector_type: str = "downsample_mlp_gelu", + input_dim: int = 1152, + n_embed: int = 2048, + depth: int = 2, + mlp_ratio: int = 1, + downsample_ratio: int = 2, + **kwargs, + ): + self.projector_type = projector_type + self.input_dim = input_dim + self.n_embed = n_embed + self.depth = depth + self.mlp_ratio = mlp_ratio + self.downsample_ratio = downsample_ratio + + super().__init__(**kwargs) + + +class DeepseekV2Config(PretrainedConfig): + + model_type = "deepseek_v2" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=102400, + hidden_size=4096, + intermediate_size=11008, + moe_intermediate_size=1407, + num_hidden_layers=30, + num_attention_heads=32, + num_key_value_heads=32, + n_shared_experts=None, + n_routed_experts=None, + ep_size=1, + routed_scaling_factor=1.0, + kv_lora_rank=512, + q_lora_rank=1536, + qk_rope_head_dim=64, + v_head_dim=128, + qk_nope_head_dim=128, + topk_method="gready", + n_group=None, + topk_group=None, + num_experts_per_tok=None, + moe_layer_freq=1, + first_k_dense_replace=0, + norm_topk_prob=False, + scoring_func="softmax", + aux_loss_alpha=0.001, + seq_aux=True, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=100000, + eos_token_id=100001, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + use_mla=True, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.ep_size = ep_size + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.topk_method = topk_method + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.moe_layer_freq = moe_layer_freq + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + self.aux_loss_alpha = aux_loss_alpha + self.seq_aux = seq_aux + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = float(rms_norm_eps) + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.use_mla = use_mla + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class DeepseekVL2Config(PretrainedConfig): + model_type = "deepseek_vl_v2" + vision_config: DeepseekVL2VisionEncoderConfig + projector_config: DeepseekVL2MlpProjectorConfig + language_config: DeepseekV2Config + + tile_tag: str = "2D" + global_view_pos: str = "head" + candidate_resolutions: Tuple[Tuple[int, int]] = ((384, 384),) + + def __init__( + self, + tile_tag: str = "tile_tag", + global_view_pos: str = "head", + candidate_resolutions: Tuple[Tuple[int, int]] = ((384, 384),), + **kwargs, + ): + super().__init__(**kwargs) + + vision_config = kwargs.get("vision_config", {}) + self.vision_config = DeepseekVL2VisionEncoderConfig(**vision_config) + + projector_config = kwargs.get("projector_config", {}) + self.projector_config = DeepseekVL2MlpProjectorConfig(**projector_config) + + language_config = kwargs.get("language_config", {}) + if isinstance(language_config, DeepseekV2Config): + self.language_config = language_config + else: + self.language_config = DeepseekV2Config(**language_config) + + self.tile_tag = tile_tag + self.global_view_pos = global_view_pos + self.candidate_resolutions = candidate_resolutions + self.architectures = ["DeepseekVL2ForCausalLM"] + + +AutoProcessor.register(DeepseekVL2Config, DeepseekVLV2Processor) diff --git a/python/sglang/srt/configs/device_config.py b/python/sglang/srt/configs/device_config.py new file mode 100644 index 000000000..3b9d3a1ed --- /dev/null +++ b/python/sglang/srt/configs/device_config.py @@ -0,0 +1,17 @@ +import logging +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + + +class DeviceConfig: + device: Optional[torch.device] + + def __init__(self, device: str = "cuda") -> None: + if device in ["cuda", "xpu", "hpu", "cpu", "npu"]: + self.device_type = device + else: + raise RuntimeError(f"Not supported device type: {device}") + self.device = torch.device(self.device_type) diff --git a/python/sglang/srt/configs/exaone.py b/python/sglang/srt/configs/exaone.py new file mode 100644 index 000000000..7b0a2d290 --- /dev/null +++ b/python/sglang/srt/configs/exaone.py @@ -0,0 +1,195 @@ +# coding=utf-8 +# Copyright 2024 The LG AI Research EXAONE Lab. All rights reserved. +# Copyright 2024 The LG CNS AI Engineering Team. +# Copyright 2023-2024 SGLang 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. +""" EXAONE model configuration """ +from typing import Any, Dict + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +EXAONE_PRETRAINED_CONFIG_ARCHIVE_MAP: Dict[str, Any] = {} + + +# ruff: noqa: E501 +class ExaoneConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a :class:`~transformers.ExaoneModel`. It is used to + instantiate a EXAONE model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the Exaone + + Configuration objects inherit from :class:`~transformers.PretrainedConfig` and can be used to control the model + outputs. Read the documentation from :class:`~transformers.PretrainedConfig` for more information. + + + Args: + vocab_size (:obj:`int`, `optional`, defaults to 102400): + Vocabulary size of the EXAONE model. Defines the number of different tokens that can be represented by the + :obj:`inputs_ids` passed when calling :class:`~transformers.ExaoneModel`. Vocabulary size of the model. + Defines the different tokens that can be represented by the `inputs_ids` passed to the forward method of + :class:`~transformers.EXAONEModel`. + max_position_embeddings (:obj:`int`, `optional`, defaults to 2048): + The maximum sequence length that this model might ever be used with. Typically set this to something large + just in case (e.g., 512 or 1024 or 2048). + hidden_size (:obj:`int`, `optional`, defaults to 2048): + Dimensionality of the encoder layers and the pooler layer. + num_layers (:obj:`int`, `optional`, defaults to 32): + Number of hidden layers in the Transformer encoder. + num_attention_heads (:obj:`int`, `optional`, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + num_key_value_heads (:obj:`int`, `optional`): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + intermediate_size (:obj:`int`, `optional`, defaults to `hidden_size * 4`): + Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. + activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"silu"`): + The non-linear activation function (function or string) in the decoder. + rope_theta (:obj:`float`, `optional`, defaults to 10000.0): + The base period of the RoPE embeddings. + rope_scaling (:obj:`Dict`, `optional`): + Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type + and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value + accordingly. + Expected contents: + `rope_type` (:obj:`str`): + The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', + 'llama3'], with 'default' being the original RoPE implementation. + `factor` (:obj:`float`, `optional`): + Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In + most scaling types, a `factor` of x will enable the model to handle sequences of length x * + original maximum pre-trained length. + `original_max_position_embeddings` (:obj:`int`, `optional`): + Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during + pretraining. + `attention_factor` (:obj:`float`, `optional`): + Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention + computation. If unspecified, it defaults to value recommended by the implementation, using the + `factor` field to infer the suggested value. + `beta_fast` (:obj:`float`, `optional`): + Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear + ramp function. If unspecified, it defaults to 32. + `beta_slow` (:obj:`float`, `optional`): + Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear + ramp function. If unspecified, it defaults to 1. + `short_factor` (:obj:`List[float]`, `optional`): + Only used with 'longrope'. The scaling factor to be applied to short contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `long_factor` (:obj:`List[float]`, `optional`): + Only used with 'longrope'. The scaling factor to be applied to long contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `low_freq_factor` (:obj:`float`, `optional`): + Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE + `high_freq_factor` (:obj:`float`, `optional`): + Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE + embed_dropout (:obj:`float`, `optional`, defaults to 0.0): + The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler. + attention_dropout (:obj:`float`, `optional`, defaults to 0.0): + The dropout ratio for the attention probabilities. + layer_norm_epsilon (:obj:`float`, `optional`, defaults to 1e-5): + The epsilon used by the layer normalization layers. + initializer_range (:obj:`float`, `optional`, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if ``configs.is_decoder=True``. + bos_token_id (:obj:`int`, `optional`, defaults to 0): + Beginning of stream token id. + eos_token_id (:obj:`int`, `optional`, defaults to 2): + End of stream token id. + tie_word_embeddings (:obj:`bool`, `optional`, defaults to :obj:`True`): + Whether to tie weight embeddings + gradient_checkpointing (:obj:`bool`, `optional`, defaults to :obj:`False`): + If True, use gradient checkpointing to save memory at the expense of slower backward pass. + + Example:: + + >>> from transformers import EXAONEModel, ExaoneConfig + + >>> # Initializing a EXAONE configuration + >>> configuration = ExaoneConfig() + + >>> # Initializing a model from configuration + >>> model = EXAONEModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.configs + """ + + model_type = "exaone" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"num_hidden_layers": "num_layers"} + + def __init__( + self, + vocab_size=102400, + max_position_embeddings=2048, + hidden_size=2048, + num_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + intermediate_size=None, + activation_function="silu", + rope_theta=10000.0, + rope_scaling=None, + embed_dropout=0.0, + attention_dropout=0.0, + layer_norm_epsilon=1e-5, + initializer_range=0.02, + use_cache=True, + bos_token_id=0, + eos_token_id=2, + tie_word_embeddings=True, + **kwargs + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.num_layers = num_layers + self.num_attention_heads = num_attention_heads + self.num_hidden_layers = num_layers + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + if intermediate_size: + self.intermediate_size = intermediate_size + else: + self.intermediate_size = hidden_size * 4 + self.activation_function = activation_function + self.embed_dropout = embed_dropout + self.attention_dropout = attention_dropout + self.layer_norm_epsilon = layer_norm_epsilon + self.initializer_range = initializer_range + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs + ) diff --git a/python/sglang/srt/configs/internvl.py b/python/sglang/srt/configs/internvl.py new file mode 100644 index 000000000..3ba9c61c1 --- /dev/null +++ b/python/sglang/srt/configs/internvl.py @@ -0,0 +1,706 @@ +import copy +import os +from shutil import copyfile +from typing import Any, Dict, List, Optional, Tuple, Union + +import sentencepiece as spm +from transformers import ( + TOKENIZER_MAPPING, + GptOssConfig, + LlamaConfig, + PretrainedConfig, + PreTrainedTokenizer, + Qwen2Config, + Qwen3Config, + Qwen3MoeConfig, +) + +from sglang.utils import logger + +# Copied from: https://github.com/OpenGVLab/InternVL/blob/34a81000402bf8f716bab8c9b57aff1f6b436bd0/internvl_chat/internvl/model/internvl_chat/configuration_internvl_chat.py#L21 + + +VOCAB_FILES_NAMES = {"vocab_file": "./tokenizer.model"} + +PRETRAINED_VOCAB_FILES_MAP = {} + + +# Modified from transformers.model.llama.configuration_llama.LlamaConfig +class InternLM2Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`InternLM2Model`]. It is used to instantiate + an InternLM2 model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the InternLM2-7B. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 32000): + Vocabulary size of the InternLM2 model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`InternLM2Model`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 11008): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer encoder. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer encoder. + num_key_value_heads (`int`, *optional*): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 2048): + The maximum sequence length that this model might ever be used with. Typically set this to something large + just in case (e.g., 512 or 1024 or 2048). + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-12): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + tie_word_embeddings(`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings + Example: + + """ + + model_type = "internlm2" + _auto_class = "AutoConfig" + + def __init__( # pylint: disable=W0102 + self, + vocab_size=103168, + hidden_size=4096, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + tie_word_embeddings=False, + bias=True, + rope_theta=10000, + rope_scaling=None, + attn_implementation="eager", + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.bias = bias + + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self._rope_scaling_validation() + + self.attn_implementation = attn_implementation + if self.attn_implementation is None: + self.attn_implementation = "eager" + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + def _rope_scaling_validation(self): + """ + Validate the `rope_scaling` configuration. + """ + if self.rope_scaling is None: + return + + if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2: + raise ValueError( + "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, " + f"got {self.rope_scaling}" + ) + rope_scaling_type = self.rope_scaling.get("type", None) + rope_scaling_factor = self.rope_scaling.get("factor", None) + if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]: + raise ValueError( + f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}" + ) + if ( + rope_scaling_factor is None + or not isinstance(rope_scaling_factor, (float, int)) + or rope_scaling_factor < 1.0 + ): + raise ValueError( + f"`rope_scaling`'s factor field must be a float|int >= 1, got {rope_scaling_factor=}, {type(rope_scaling_factor)=}" + ) + if isinstance(rope_scaling_factor, int): + rope_scaling_factor = float(rope_scaling_factor) + + +class InternVisionConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`InternVisionModel`]. It is used to + instantiate a vision encoder according to the specified arguments, defining the model architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + num_channels (`int`, *optional*, defaults to 3): + Number of color channels in the input images (e.g., 3 for RGB). + patch_size (`int`, *optional*, defaults to 14): + The size (resolution) of each patch. + image_size (`int`, *optional*, defaults to 224): + The size (resolution) of each image. + qkv_bias (`bool`, *optional*, defaults to `False`): + Whether to add a bias to the queries and values in the self-attention layers. + hidden_size (`int`, *optional*, defaults to 3200): + Dimensionality of the encoder layers and the pooler layer. + num_attention_heads (`int`, *optional*, defaults to 25): + Number of attention heads for each attention layer in the Transformer encoder. + intermediate_size (`int`, *optional*, defaults to 12800): + Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. + qk_normalization (`bool`, *optional*, defaults to `True`): + Whether to normalize the queries and keys in the self-attention layers. + num_hidden_layers (`int`, *optional*, defaults to 48): + Number of hidden layers in the Transformer encoder. + use_flash_attn (`bool`, *optional*, defaults to `True`): + Whether to use flash attention mechanism. + hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`): + The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, + `"relu"`, `"selu"` and `"gelu_new"` ``"gelu"` are supported. + layer_norm_eps (`float`, *optional*, defaults to 1e-6): + The epsilon used by the layer normalization layers. + dropout (`float`, *optional*, defaults to 0.0): + The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. + drop_path_rate (`float`, *optional*, defaults to 0.0): + Dropout rate for stochastic depth. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + initializer_factor (`float`, *optional*, defaults to 0.1): + A factor for layer scale. + """ + + model_type = "intern_vit_6b" + + def __init__( + self, + num_channels=3, + patch_size=14, + image_size=224, + qkv_bias=False, + hidden_size=3200, + num_attention_heads=25, + intermediate_size=12800, + qk_normalization=True, + num_hidden_layers=48, + use_flash_attn=True, + hidden_act="gelu", + layer_norm_eps=1e-6, + dropout=0.0, + drop_path_rate=0.0, + attention_dropout=0.0, + initializer_range=0.02, + initializer_factor=0.1, + **kwargs, + ): + super().__init__(**kwargs) + + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.dropout = dropout + self.drop_path_rate = drop_path_rate + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_channels = num_channels + self.patch_size = patch_size + self.image_size = image_size + self.initializer_range = initializer_range + self.initializer_factor = initializer_factor + self.attention_dropout = attention_dropout + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.qkv_bias = qkv_bias + self.qk_normalization = qk_normalization + self.use_flash_attn = use_flash_attn + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs + ) -> "PretrainedConfig": + config_dict, kwargs = cls.get_config_dict( + pretrained_model_name_or_path, **kwargs + ) + + if "vision_config" in config_dict: + config_dict = config_dict["vision_config"] + + if ( + "model_type" in config_dict + and hasattr(cls, "model_type") + and config_dict["model_type"] != cls.model_type + ): + logger.warning( + f"You are using a model of type {config_dict['model_type']} to instantiate a model of type " + f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." + ) + + return cls.from_dict(config_dict, **kwargs) + + +class InternVLChatConfig(PretrainedConfig): + model_type = "internvl_chat" + is_composition = True + + def __init__( + self, + vision_config=None, + llm_config=None, + use_backbone_lora=0, + use_llm_lora=0, + pad2square=False, + select_layer=-1, + force_image_size=None, + downsample_ratio=0.5, + template=None, + dynamic_image_size=False, + use_thumbnail=False, + ps_version="v1", + min_dynamic_patch=1, + max_dynamic_patch=6, + **kwargs, + ): + super().__init__(**kwargs) + + if vision_config is None: + vision_config = {"architectures": ["InternVisionModel"]} + logger.info( + "vision_config is None. Initializing the InternVisionConfig with default values." + ) + + if llm_config is None: + llm_config = {"architectures": ["InternLM2ForCausalLM"]} + logger.info( + "llm_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`)." + ) + + self.vision_config = InternVisionConfig(**vision_config) + if llm_config.get("architectures")[0] == "LlamaForCausalLM": + self.llm_config = LlamaConfig(**llm_config) + elif llm_config.get("architectures")[0] == "InternLM2ForCausalLM": + self.llm_config = InternLM2Config(**llm_config) + elif llm_config.get("architectures")[0] == "Qwen2ForCausalLM": + self.llm_config = Qwen2Config(**llm_config) + elif llm_config.get("architectures")[0] == "Qwen3MoeForCausalLM": + self.llm_config = Qwen3MoeConfig(**llm_config) + elif llm_config.get("architectures")[0] == "Qwen3ForCausalLM": + self.llm_config = Qwen3Config(**llm_config) + elif llm_config.get("architectures")[0] == "GptOssForCausalLM": + self.llm_config = GptOssConfig(**llm_config) + else: + raise ValueError( + "Unsupported architecture: {}".format( + llm_config.get("architectures")[0] + ) + ) + + self.use_backbone_lora = use_backbone_lora + self.use_llm_lora = use_llm_lora + self.pad2square = pad2square + self.select_layer = select_layer + self.force_image_size = force_image_size + self.downsample_ratio = downsample_ratio + self.template = template + self.dynamic_image_size = dynamic_image_size + self.use_thumbnail = use_thumbnail + self.ps_version = ps_version # pixel shuffle version + self.min_dynamic_patch = min_dynamic_patch + self.max_dynamic_patch = max_dynamic_patch + + self.hidden_size = self.llm_config.hidden_size + # By default, we use tie_word_embeddings=False for models of all sizes. + self.tie_word_embeddings = False + self.llm_config.tie_word_embeddings = self.tie_word_embeddings + + def to_dict(self): + """ + Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. + + Returns: + `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance, + """ + output = copy.deepcopy(self.__dict__) + output["vision_config"] = self.vision_config.to_dict() + output["llm_config"] = self.llm_config.to_dict() + output["model_type"] = self.__class__.model_type + output["use_backbone_lora"] = self.use_backbone_lora + output["use_llm_lora"] = self.use_llm_lora + output["select_layer"] = self.select_layer + output["force_image_size"] = self.force_image_size + output["downsample_ratio"] = self.downsample_ratio + output["template"] = self.template + output["dynamic_image_size"] = self.dynamic_image_size + output["use_thumbnail"] = self.use_thumbnail + output["ps_version"] = self.ps_version + output["min_dynamic_patch"] = self.min_dynamic_patch + output["max_dynamic_patch"] = self.max_dynamic_patch + + return output + + +# # Modified from transformers.model.llama.tokenization_llama_fast.LlamaTokenizerFast -> InternLM2TokenizerFast +# class InternLM2TokenizerFast(PreTrainedTokenizerFast): +# vocab_files_names = VOCAB_FILES_NAMES +# slow_tokenizer_class = InternLM2Tokenizer +# padding_side = 'left' +# model_input_names = ['input_ids', 'attention_mask'] +# _auto_class = 'AutoTokenizer' +# +# def __init__( +# self, +# vocab_file, +# unk_token='', +# bos_token='', +# eos_token='', +# pad_token='', +# sp_model_kwargs: Optional[Dict[str, Any]] = None, +# add_bos_token=True, +# add_eos_token=False, +# decode_with_prefix_space=False, +# clean_up_tokenization_spaces=False, +# **kwargs, +# ): +# super().__init__( +# vocab_file=vocab_file, +# unk_token=unk_token, +# bos_token=bos_token, +# eos_token=eos_token, +# pad_token=pad_token, +# sp_model_kwargs=sp_model_kwargs, +# add_bos_token=add_bos_token, +# add_eos_token=add_eos_token, +# decode_with_prefix_space=decode_with_prefix_space, +# clean_up_tokenization_spaces=clean_up_tokenization_spaces, +# **kwargs, +# ) +# self._add_bos_token = add_bos_token +# self._add_eos_token = add_eos_token +# self.update_post_processor() +# self.vocab_file = vocab_file +# +# @property +# def can_save_slow_tokenizer(self) -> bool: +# return os.path.isfile(self.vocab_file) if self.vocab_file else False +# +# def update_post_processor(self): +# """ +# Updates the underlying post processor with the current `bos_token` and `eos_token`. +# """ +# bos = self.bos_token +# bos_token_id = self.bos_token_id +# if bos is None and self.add_bos_token: +# raise ValueError('add_bos_token = True but bos_token = None') +# +# eos = self.eos_token +# eos_token_id = self.eos_token_id +# if eos is None and self.add_eos_token: +# raise ValueError('add_eos_token = True but eos_token = None') +# +# single = f"{(bos + ':0 ') if self.add_bos_token else ''}$A:0{(' ' + eos + ':0') if self.add_eos_token else ''}" +# pair = f"{single}{(' ' + bos + ':1') if self.add_bos_token else ''} $B:1{(' ' + eos + ':1') if self.add_eos_token else ''}" +# +# special_tokens = [] +# if self.add_bos_token: +# special_tokens.append((bos, bos_token_id)) +# if self.add_eos_token: +# special_tokens.append((eos, eos_token_id)) +# self._tokenizer.post_processor = processors.TemplateProcessing( +# single=single, pair=pair, special_tokens=special_tokens +# ) +# +# @property +# def add_eos_token(self): +# return self._add_eos_token +# +# @property +# def add_bos_token(self): +# return self._add_bos_token +# +# @add_eos_token.setter +# def add_eos_token(self, value): +# self._add_eos_token = value +# self.update_post_processor() +# +# @add_bos_token.setter +# def add_bos_token(self, value): +# self._add_bos_token = value +# self.update_post_processor() +# +# def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]: +# if not self.can_save_slow_tokenizer: +# raise ValueError( +# 'Your fast tokenizer does not have the necessary information to save the vocabulary for a slow ' +# 'tokenizer.' +# ) +# +# if not os.path.isdir(save_directory): +# logger.error(f'Vocabulary path ({save_directory}) should be a directory') +# return +# out_vocab_file = os.path.join( +# save_directory, (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'] +# ) +# +# if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file): +# copyfile(self.vocab_file, out_vocab_file) +# +# return (out_vocab_file,) + + +# Modified from transformers.model.llama.tokenization_llama.LlamaTokenizer +class InternLM2Tokenizer(PreTrainedTokenizer): + """ + Construct a InternLM2 tokenizer. Based on byte-level Byte-Pair-Encoding. + + Args: + vocab_file (`str`): + Path to the vocabulary file. + """ + + vocab_files_names = VOCAB_FILES_NAMES + pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP + model_input_names = ["input_ids", "attention_mask"] + _auto_class = "AutoTokenizer" + + def __init__( + self, + vocab_file, + unk_token="", + bos_token="", + eos_token="", + pad_token="", + sp_model_kwargs: Optional[Dict[str, Any]] = None, + add_bos_token=True, + add_eos_token=False, + decode_with_prefix_space=False, + clean_up_tokenization_spaces=False, + **kwargs, + ): + print("register succeed") + self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs + self.vocab_file = vocab_file + self.add_bos_token = add_bos_token + self.add_eos_token = add_eos_token + self.decode_with_prefix_space = decode_with_prefix_space + self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs) + self.sp_model.Load(vocab_file) + self._no_prefix_space_tokens = None + super().__init__( + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + @property + def no_prefix_space_tokens(self): + if self._no_prefix_space_tokens is None: + vocab = self.convert_ids_to_tokens(list(range(self.vocab_size))) + self._no_prefix_space_tokens = { + i for i, tok in enumerate(vocab) if not tok.startswith("▁") + } + return self._no_prefix_space_tokens + + @property + def vocab_size(self): + """Returns vocab size""" + return self.sp_model.get_piece_size() + + @property + def bos_token_id(self) -> Optional[int]: + return self.sp_model.bos_id() + + @property + def eos_token_id(self) -> Optional[int]: + return self.sp_model.eos_id() + + def get_vocab(self): + """Returns vocab as a dict""" + vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)} + vocab.update(self.added_tokens_encoder) + return vocab + + def _tokenize(self, text): + """Returns a tokenized string.""" + return self.sp_model.encode(text, out_type=str) + + def _convert_token_to_id(self, token): + """Converts a token (str) in an id using the vocab.""" + return self.sp_model.piece_to_id(token) + + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + token = self.sp_model.IdToPiece(index) + return token + + def _maybe_add_prefix_space(self, tokens, decoded): + if tokens and tokens[0] not in self.no_prefix_space_tokens: + return " " + decoded + else: + return decoded + + def convert_tokens_to_string(self, tokens): + """Converts a sequence of tokens (string) in a single string.""" + current_sub_tokens = [] + out_string = "" + prev_is_special = False + for token in tokens: + # make sure that special tokens are not decoded using sentencepiece model + if token in self.all_special_tokens: + if not prev_is_special: + out_string += " " + out_string += self.sp_model.decode(current_sub_tokens) + token + prev_is_special = True + current_sub_tokens = [] + else: + current_sub_tokens.append(token) + prev_is_special = False + out_string += self.sp_model.decode(current_sub_tokens) + out_string = self.clean_up_tokenization(out_string) + out_string = self._maybe_add_prefix_space(tokens=tokens, decoded=out_string) + return out_string[1:] + + def save_vocabulary( + self, save_directory, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + """ + Save the vocabulary and special tokens file to a directory. + + Args: + save_directory (`str`): + The directory in which to save the vocabulary. + + Returns: + `Tuple(str)`: Paths to the files saved. + """ + if not os.path.isdir(save_directory): + logger.error(f"Vocabulary path ({save_directory}) should be a directory") + return + out_vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"], + ) + + if os.path.abspath(self.vocab_file) != os.path.abspath( + out_vocab_file + ) and os.path.isfile(self.vocab_file): + copyfile(self.vocab_file, out_vocab_file) + elif not os.path.isfile(self.vocab_file): + with open(out_vocab_file, "wb") as fi: + content_spiece_model = self.sp_model.serialized_model_proto() + fi.write(content_spiece_model) + + return (out_vocab_file,) + + def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): + if self.add_bos_token: + bos_token_ids = [self.bos_token_id] + else: + bos_token_ids = [] + + output = bos_token_ids + token_ids_0 + + if token_ids_1 is not None: + output = output + token_ids_1 + + if self.add_eos_token: + output = output + [self.eos_token_id] + + return output + + def get_special_tokens_mask( + self, + token_ids_0: List[int], + token_ids_1: Optional[List[int]] = None, + already_has_special_tokens: bool = False, + ) -> List[int]: + """ + Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding + special tokens using the tokenizer `prepare_for_model` method. + + Args: + token_ids_0 (`List[int]`): + List of IDs. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + if already_has_special_tokens: + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, + token_ids_1=token_ids_1, + already_has_special_tokens=True, + ) + + if token_ids_1 is None: + return [1] + ([0] * len(token_ids_0)) + [1] + return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1] + + def create_token_type_ids_from_sequences( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """ + Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make + use of token type ids, therefore a list of zeros is returned. + + Args: + token_ids_0 (`List[int]`): + List of IDs. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `List[int]`: List of zeros. + """ + eos = [self.eos_token_id] + + if token_ids_1 is None: + return len(token_ids_0 + eos) * [0] + return len(token_ids_0 + eos + token_ids_1 + eos) * [0] + + +TOKENIZER_MAPPING.register( + InternVLChatConfig, (InternLM2Tokenizer, None), exist_ok=True +) diff --git a/python/sglang/srt/configs/janus_pro.py b/python/sglang/srt/configs/janus_pro.py new file mode 100644 index 000000000..d574953e9 --- /dev/null +++ b/python/sglang/srt/configs/janus_pro.py @@ -0,0 +1,634 @@ +# Adapted from: +# https://github.com/deepseek-ai/Janus/tree/main/janus/models + +from dataclasses import dataclass +from typing import Dict, List, Tuple, Union + +import numpy as np +import PIL +import torch +from PIL.Image import Image +from transformers import ( + BaseImageProcessor, + BatchFeature, + LlamaConfig, + LlamaTokenizerFast, + PretrainedConfig, + ProcessorMixin, +) +from transformers.image_utils import to_numpy_array + +from sglang.srt.configs.utils import register_image_processor, register_processor +from sglang.srt.multimodal.mm_utils import expand2square + + +class DictToObject(dict): + def __init__(self, dictionary): + super(self).__init__(dictionary) + + for key, value in dictionary.items(): + if isinstance(value, dict): + value = DictToObject(value) + setattr(self, key, value) + + +class VisionConfig(PretrainedConfig): + model_type = "vision" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +class GenAlignerConfig(PretrainedConfig): + model_type = "gen_aligner" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +class GenHeadConfig(PretrainedConfig): + model_type = "gen_head" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +class AlignerConfig(PretrainedConfig): + model_type = "aligner" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +class GenVisionConfig(PretrainedConfig): + model_type = "gen_vision" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +@dataclass +class SigLIPVisionCfg: + width: int = 1152 + layers: Union[Tuple[int, int, int, int], int] = 27 + heads: int = 16 + patch_size: int = 14 + image_size: Union[Tuple[int, int], int] = 336 + global_pool: str = "map" + mlp_ratio: float = 3.7362 + class_token: bool = False + num_classes: int = 0 + use_checkpoint: bool = False + + +class MultiModalityConfig(PretrainedConfig): + model_type = "multi_modality" + vision_config: VisionConfig + aligner_config: AlignerConfig + + gen_vision_config: GenVisionConfig + gen_aligner_config: GenAlignerConfig + gen_head_config: GenHeadConfig + + language_config: LlamaConfig + + def __init__(self, **kwargs): + super().__init__(**kwargs) + vision_config = kwargs.get("vision_config", {}) + self.vision_config = VisionConfig(**vision_config) + + aligner_config = kwargs.get("aligner_config", {}) + self.aligner_config = AlignerConfig(**aligner_config) + + gen_vision_config = kwargs.get("gen_vision_config", {}) + self.gen_vision_config = GenVisionConfig(**gen_vision_config) + + gen_aligner_config = kwargs.get("gen_aligner_config", {}) + self.gen_aligner_config = GenAlignerConfig(**gen_aligner_config) + + gen_head_config = kwargs.get("gen_head_config", {}) + self.gen_head_config = GenHeadConfig(**gen_head_config) + + language_config = kwargs.get("language_config", {}) + if isinstance(language_config, LlamaConfig): + self.language_config = language_config + else: + self.language_config = LlamaConfig(**language_config) + + +class VLMImageProcessor(BaseImageProcessor): + model_input_names = ["pixel_values"] + + def __init__( + self, + image_size: int, + min_size: int = 14, + image_mean: Union[Tuple[float, float, float], List[float]] = ( + 0.48145466, + 0.4578275, + 0.40821073, + ), + image_std: Union[Tuple[float, float, float], List[float]] = ( + 0.26862954, + 0.26130258, + 0.27577711, + ), + rescale_factor: float = 1.0 / 255.0, + do_normalize: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + + self.image_size = image_size + self.rescale_factor = rescale_factor + self.image_mean = image_mean + self.image_std = image_std + self.min_size = min_size + self.do_normalize = do_normalize + + if image_mean is None: + self.background_color = (127, 127, 127) + else: + self.background_color = tuple([int(x * 255) for x in image_mean]) + + def resize(self, pil_img: Image) -> np.ndarray: + """ + + Args: + pil_img (PIL.Image): [H, W, 3] in PIL.Image in RGB + + Returns: + x (np.ndarray): [3, self.image_size, self.image_size] + """ + + width, height = pil_img.size + max_size = max(width, height) + + size = [ + max(int(height / max_size * self.image_size), self.min_size), + max(int(width / max_size * self.image_size), self.min_size), + ] + + if width <= 0 or height <= 0 or size[0] <= 0 or size[1] <= 0: + # print(f"orig size = {pil_img.size}, new size = {size}") + raise ValueError("Invalid size!") + + def resize( + pil_img, size, interpolation=PIL.Image.Resampling.BICUBIC, antialias=True + ): + if isinstance(size, int): + w, h = pil_img.size + if (w <= h and w == size) or (h <= w and h == size): + return pil_img + if w < h: + ow = size + oh = int(size * h / w) + else: + oh = size + ow = int(size * w / h) + size = (ow, oh) + else: + size = (size[1], size[0]) + + return pil_img.resize( + size, resample=interpolation, reducing_gap=None if antialias else 3.0 + ) + + pil_img = resize( + pil_img, size, interpolation=PIL.Image.Resampling.BICUBIC, antialias=True + ) + + pil_img = expand2square(pil_img, self.background_color) + x = to_numpy_array(pil_img) + + # [H, W, 3] -> [3, H, W] + x = np.transpose(x, (2, 0, 1)) + + return x + + def preprocess(self, images, return_tensors: str = "pt", **kwargs) -> BatchFeature: + # resize and pad to [self.image_size, self.image_size] + # then convert from [H, W, 3] to [3, H, W] + if not isinstance(images, list): + images = [images] + images: List[np.ndarray] = [self.resize(image) for image in images] + images = [image[:3, ...] for image in images] + + # rescale from [0, 255] -> [0, 1] + images = [ + self.rescale( + image=image, + scale=self.rescale_factor, + input_data_format="channels_first", + ) + for image in images + ] + + # normalize + if self.do_normalize: + images = [ + self.normalize( + image=image, + mean=self.image_mean, + std=self.image_std, + input_data_format="channels_first", + ) + for image in images + ] + data = {"pixel_values": images} + return BatchFeature(data=data, tensor_type=return_tensors) + + @property + def default_shape(self): + return [3, self.image_size, self.image_size] + + +class DictOutput(object): + def items(self): + return self.__dict__.items() + + def keys(self): + return self.__dict__.keys() + + def __getitem__(self, item): + return self.__dict__[item] + + def __contains__(self, key): + return key in self.__dict__ + + def __setitem__(self, key, value): + self.__dict__[key] = value + + +@dataclass +class VLChatProcessorOutput(DictOutput): + sft_format: str + input_ids: torch.Tensor + pixel_values: torch.Tensor + num_image_tokens: torch.IntTensor + + def __len__(self): + return len(self.input_ids) + + +@dataclass +class BatchedVLChatProcessorOutput(DictOutput): + sft_format: List[str] + input_ids: torch.Tensor + pixel_values: torch.Tensor + attention_mask: torch.Tensor + images_seq_mask: torch.BoolTensor + images_emb_mask: torch.BoolTensor + + +# FIXME: had to place Official Processor here, since image_processor module would not be imported in all threads, +# hence AutoProcessor registration would not be affective in some cases +class VLChatProcessor(ProcessorMixin): + image_processor_class = "AutoImageProcessor" + tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") + + attributes = ["image_processor", "tokenizer"] + + def __init__( + self, + image_processor: VLMImageProcessor, + tokenizer: LlamaTokenizerFast, + image_tag: str = "", + image_start_tag: str = "", + image_end_tag: str = "", + pad_tag: str = "<｜▁pad▁｜>", + num_image_tokens: int = 576, + add_special_token: bool = False, + sft_format: str = "deepseek", + mask_prompt: bool = True, + ignore_id: int = -100, + **kwargs, + ): + self.image_processor = image_processor + self.tokenizer = tokenizer + + image_id = self.tokenizer.vocab.get(image_tag) + if image_id is None: + special_tokens = [image_tag] + special_tokens_dict = {"additional_special_tokens": special_tokens} + self.tokenizer.add_special_tokens(special_tokens_dict) + # print(f"Add image tag = {image_tag} to the tokenizer") + + self.image_tag = image_tag + self.image_start_tag = image_start_tag + self.image_end_tag = image_end_tag + self.pad_tag = pad_tag + + self.num_image_tokens = num_image_tokens + self.add_special_token = add_special_token + self.sft_format = sft_format + self.ignore_id = ignore_id + + super().__init__( + image_processor, + tokenizer, + **kwargs, + ) + + @property + def image_token(self): + return self.image_tag + + @property + def image_id(self) -> int: + image_id = self.tokenizer.vocab.get(self.image_tag) + return image_id + + @property + def image_start_id(self): + image_start_id = self.tokenizer.vocab.get(self.image_start_tag) + return image_start_id + + @property + def image_end_id(self): + image_end_id = self.tokenizer.vocab.get(self.image_end_tag) + return image_end_id + + @property + def image_start_token(self): + return self.image_start_tag + + @property + def image_end_token(self): + return self.image_end_tag + + @property + def pad_id(self): + pad_id = self.tokenizer.vocab.get(self.pad_tag) + return pad_id + + def add_image_token( + self, + image_indices: List[int], + input_ids: torch.LongTensor, + ): + """ + + Args: + image_indices (List[int]): [index_0, index_1, ..., index_j] + input_ids (torch.LongTensor): [N] + + Returns: + input_ids (torch.LongTensor): [N + image tokens] + num_image_tokens (torch.IntTensor): [n_images] + """ + + input_slices = [] + + start = 0 + for index in image_indices: + if self.add_special_token: + end = index + 1 + else: + end = index + + # original text tokens + input_slices.append(input_ids[start:end]) + + # add boi, image tokens, eoi and set the mask as False + input_slices.append(self.image_start_id * torch.ones((1), dtype=torch.long)) + input_slices.append( + self.image_id * torch.ones((self.num_image_tokens,), dtype=torch.long) + ) + input_slices.append(self.image_end_id * torch.ones((1), dtype=torch.long)) + start = index + 1 + + # the left part + input_slices.append(input_ids[start:]) + + # concat all slices + input_ids = torch.cat(input_slices, dim=0) + num_image_tokens = torch.IntTensor([self.num_image_tokens] * len(image_indices)) + + return input_ids, num_image_tokens + + def process_one( + self, + prompt: str = None, + images: List[Image] = None, + **kwargs, + ): + """ + + Args: + prompt (str): the formatted prompt; + images (List[ImageType]): the list of images; + **kwargs: + + Returns: + outputs (BaseProcessorOutput): the output of the processor, + - input_ids (torch.LongTensor): [N + image tokens] + - target_ids (torch.LongTensor): [N + image tokens] + - images (torch.FloatTensor): [n_images, 3, H, W] + - image_id (int): the id of the image token + - num_image_tokens (List[int]): the number of image tokens + """ + + sft_format = prompt + # tokenize + input_ids = self.tokenizer.encode(sft_format) + input_ids = torch.LongTensor(input_ids) + + # add image tokens to the input_ids + image_token_mask: torch.Tensor = (input_ids == self.image_id).to(torch.bool) + image_indices = image_token_mask.nonzero() + input_ids, num_image_tokens = self.add_image_token( + image_indices=image_indices, + input_ids=input_ids, + ) + + # load images + images_outputs = self.image_processor(images, return_tensors="pt") + + prepare = VLChatProcessorOutput( + sft_format=sft_format, + input_ids=input_ids, + pixel_values=images_outputs.pixel_values, + num_image_tokens=num_image_tokens, + ) + + return prepare + + def __call__( + self, + *, + prompt: str = None, + conversations: List[Dict[str, str]] = None, + images: List[Image] = None, + force_batchify: bool = True, + **kwargs, + ): + """ + + Args: + prompt (str): the formatted prompt; + conversations (List[Dict]): conversations with a list of messages; + images (List[ImageType]): the list of images; + force_batchify (bool): force batchify the inputs; + **kwargs: + + Returns: + outputs (BaseProcessorOutput): the output of the processor, + - input_ids (torch.LongTensor): [N + image tokens] + - images (torch.FloatTensor): [n_images, 3, H, W] + - image_id (int): the id of the image token + - num_image_tokens (List[int]): the number of image tokens + """ + + prepare = self.process_one( + prompt=prompt, conversations=conversations, images=images + ) + + if force_batchify: + prepare = self.batchify([prepare]) + + return prepare + + def batchify( + self, prepare_list: List[VLChatProcessorOutput] + ) -> BatchedVLChatProcessorOutput: + """ + Preprocesses the inputs for multimodal inference. + + Args: + prepare_list (List[VLChatProcessorOutput]): A list of VLChatProcessorOutput. + + Returns: + BatchedVLChatProcessorOutput: A dictionary of the inputs to use for multimodal inference. + """ + + batch_size = len(prepare_list) + sft_format = [] + n_images = [] + seq_lens = [] + for prepare in prepare_list: + n_images.append(len(prepare.num_image_tokens)) + seq_lens.append(len(prepare)) + + input_token_max_len = max(seq_lens) + max_n_images = max(1, max(n_images)) + + batched_input_ids = torch.full( + (batch_size, input_token_max_len), self.pad_id + ).long() # FIXME + batched_attention_mask = torch.zeros((batch_size, input_token_max_len)).long() + batched_pixel_values = torch.zeros( + (batch_size, max_n_images, *self.image_processor.default_shape) + ).float() + batched_images_seq_mask = torch.zeros((batch_size, input_token_max_len)).bool() + batched_images_emb_mask = torch.zeros( + (batch_size, max_n_images, self.num_image_tokens) + ).bool() + + for i, prepare in enumerate(prepare_list): + input_ids = prepare.input_ids + seq_len = len(prepare) + n_image = len(prepare.num_image_tokens) + # left-padding + batched_attention_mask[i, -seq_len:] = 1 + batched_input_ids[i, -seq_len:] = torch.LongTensor(input_ids) + batched_images_seq_mask[i, -seq_len:] = input_ids == self.image_id + + if n_image > 0: + batched_pixel_values[i, :n_image] = prepare.pixel_values + for j, n_image_tokens in enumerate(prepare.num_image_tokens): + batched_images_emb_mask[i, j, :n_image_tokens] = True + + sft_format.append(prepare.sft_format) + + batched_prepares = BatchedVLChatProcessorOutput( + input_ids=batched_input_ids, + attention_mask=batched_attention_mask, + pixel_values=batched_pixel_values, + images_seq_mask=batched_images_seq_mask, + images_emb_mask=batched_images_emb_mask, + sft_format=sft_format, + ) + + return batched_prepares + + +class VLMImageProcessorConfig(PretrainedConfig): + model_type = "deepseek_vlm" + image_size: int + min_size: int + image_mean: Union[Tuple[float, float, float], List[float]] + image_std: Union[Tuple[float, float, float], List[float]] + rescale_factor: float + do_normalize: bool + + def __init__( + self, + image_size: int, + min_size: int = 14, + image_mean: Union[Tuple[float, float, float], List[float]] = ( + 0.48145466, + 0.4578275, + 0.40821073, + ), + image_std: Union[Tuple[float, float, float], List[float]] = ( + 0.26862954, + 0.26130258, + 0.27577711, + ), + rescale_factor: float = 1.0 / 255.0, + do_normalize: bool = True, + **kwargs, + ): + self.image_size = image_size + self.min_size = min_size + self.image_mean = image_mean + self.image_std = image_std + self.rescale_factor = rescale_factor + self.do_normalize = do_normalize + + super().__init__(**kwargs) + + +register_processor(MultiModalityConfig, VLChatProcessor) +register_image_processor(MultiModalityConfig, VLMImageProcessor) diff --git a/python/sglang/srt/configs/kimi_vl.py b/python/sglang/srt/configs/kimi_vl.py new file mode 100644 index 000000000..3c7d20f59 --- /dev/null +++ b/python/sglang/srt/configs/kimi_vl.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py +from typing import Optional, Union + +from transformers.configuration_utils import PretrainedConfig + +from sglang.srt.configs.deepseekvl2 import DeepseekV2Config +from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig + + +class KimiVLConfig(PretrainedConfig): + model_type = "kimi_vl" + + def __init__( + self, + vision_config: Optional[Union[dict, MoonViTConfig]] = None, + text_config: Optional[Union[dict, DeepseekV2Config]] = None, + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + **kwargs + ): + if vision_config is None: + vision_config = MoonViTConfig() + elif isinstance(vision_config, dict): + vision_config = MoonViTConfig(**vision_config) + self.vision_config = vision_config + + if text_config is None: + text_config = DeepseekV2Config() + elif isinstance(text_config, dict): + text_config = DeepseekV2Config(**text_config) + self.text_config = text_config + + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + + super().__init__(pad_token_id=pad_token_id, **kwargs) diff --git a/python/sglang/srt/configs/kimi_vl_moonvit.py b/python/sglang/srt/configs/kimi_vl_moonvit.py new file mode 100644 index 000000000..166809eb6 --- /dev/null +++ b/python/sglang/srt/configs/kimi_vl_moonvit.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py +from transformers.configuration_utils import PretrainedConfig + + +class MoonViTConfig(PretrainedConfig): + model_type = "moonvit" + + def __init__( + self, + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + num_attention_heads: int = 16, + num_hidden_layers: int = 27, + hidden_size: int = 1152, + intermediate_size: int = 4304, + merge_kernel_size: tuple[int, int] = (2, 2), + **kwargs, + ): + super().__init__(**kwargs) + self.patch_size = patch_size + # Positional embedding config + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + # Transformer config + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + # Patch merger config + self.merge_kernel_size = merge_kernel_size diff --git a/python/sglang/srt/configs/load_config.py b/python/sglang/srt/configs/load_config.py new file mode 100644 index 000000000..be9a40b4b --- /dev/null +++ b/python/sglang/srt/configs/load_config.py @@ -0,0 +1,89 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py +import enum +import json +import logging +from dataclasses import dataclass, field +from typing import List, Optional, Union + +from sglang.srt.utils import is_hip + +logger = logging.getLogger(__name__) + + +class LoadFormat(str, enum.Enum): + AUTO = "auto" + PT = "pt" + SAFETENSORS = "safetensors" + NPCACHE = "npcache" + DUMMY = "dummy" + SHARDED_STATE = "sharded_state" + GGUF = "gguf" + BITSANDBYTES = "bitsandbytes" + MISTRAL = "mistral" + LAYERED = "layered" + JAX = "jax" + REMOTE = "remote" + + +@dataclass +class LoadConfig: + """ + download_dir: Directory to download and load the weights, default to the + default cache directory of huggingface. + load_format: The format of the model weights to load: + "auto" will try to load the weights in the safetensors format and + fall back to the pytorch bin format if safetensors format is + not available. + "pt" will load the weights in the pytorch bin format. + "safetensors" will load the weights in the safetensors format. + "npcache" will load the weights in pytorch format and store + a numpy cache to speed up the loading. + "dummy" will initialize the weights with random values, which is + mainly for profiling. + "bitsandbytes" will load nf4 type weights. + ignore_patterns: The list of patterns to ignore when loading the model. + Default to "original/**/*" to avoid repeated loading of llama's + checkpoints. + decryption_key_file: If set, decrypts the output files with a password read + from this file (after PBKDF2). + """ + + load_format: Union[str, LoadFormat] = LoadFormat.AUTO + download_dir: Optional[str] = None + model_loader_extra_config: Optional[Union[str, dict]] = field(default_factory=dict) + ignore_patterns: Optional[Union[List[str], str]] = None + decryption_key_file: Optional[str] = None + + def __post_init__(self): + model_loader_extra_config = self.model_loader_extra_config or {} + if isinstance(model_loader_extra_config, str): + self.model_loader_extra_config = json.loads(model_loader_extra_config) + self._verify_load_format() + + if self.ignore_patterns is not None and len(self.ignore_patterns) > 0: + logger.info( + "Ignoring the following patterns when downloading weights: %s", + self.ignore_patterns, + ) + else: + self.ignore_patterns = ["original/**/*"] + + def _verify_load_format(self) -> None: + if not isinstance(self.load_format, str): + return + + load_format = self.load_format.lower() + self.load_format = LoadFormat(load_format) + + rocm_not_supported_load_format: List[str] = [] + if is_hip() and load_format in rocm_not_supported_load_format: + rocm_supported_load_format = [ + f + for f in LoadFormat.__members__ + if (f not in rocm_not_supported_load_format) + ] + raise ValueError( + f"load format '{load_format}' is not supported in ROCm. " + f"Supported load formats are " + f"{rocm_supported_load_format}" + ) diff --git a/python/sglang/srt/configs/longcat_flash.py b/python/sglang/srt/configs/longcat_flash.py new file mode 100644 index 000000000..e6a2dfb02 --- /dev/null +++ b/python/sglang/srt/configs/longcat_flash.py @@ -0,0 +1,104 @@ +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +FLASH_PRETRAINED_CONFIG_ARCHIVE_MAP = {} + + +class LongcatFlashConfig(PretrainedConfig): + model_type = "longcat_flash" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=131072, + hidden_size=6144, + intermediate_size=None, + ffn_hidden_size=12288, + expert_ffn_hidden_size=2048, + num_layers=28, + num_hidden_layers=None, + num_attention_heads=64, + ep_size=1, + kv_lora_rank=512, + q_lora_rank=1536, + qk_rope_head_dim=128, + qk_nope_head_dim=128, + v_head_dim=128, + n_routed_experts=512, + moe_topk=12, + norm_topk_prob=False, + max_position_embeddings=131072, + rms_norm_eps=1e-05, + use_cache=True, + pad_token_id=None, + bos_token_id=1, + eos_token_id=2, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + mla_scale_q_lora=True, + mla_scale_kv_lora=True, + torch_dtype="bfloat16", + params_dtype="bfloat16", + rounter_params_dtype="float32", + router_bias=False, + topk_method=None, + routed_scaling_factor=6.0, + zero_expert_num=256, + zero_expert_type="identity", + nextn_use_scmoe=False, + num_nextn_predict_layers=1, + **kwargs, + ): + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + torch_dtype=torch_dtype, + params_dtype=params_dtype, + rounter_params_dtype=rounter_params_dtype, + topk_method=topk_method, + router_bias=router_bias, + nextn_use_scmoe=nextn_use_scmoe, + num_nextn_predict_layers=num_nextn_predict_layers, + **kwargs, + ) + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.num_hidden_layers = ( + num_hidden_layers if num_hidden_layers is not None else num_layers + ) + self.intermediate_size = ( + intermediate_size if intermediate_size is not None else ffn_hidden_size + ) + self.moe_intermediate_size = expert_ffn_hidden_size + self.num_attention_heads = num_attention_heads + self.ep_size = ep_size + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.n_routed_experts = n_routed_experts + self.moe_topk = moe_topk + self.norm_topk_prob = norm_topk_prob + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.mla_scale_q_lora = mla_scale_q_lora + self.mla_scale_kv_lora = mla_scale_kv_lora + self.zero_expert_num = zero_expert_num + self.zero_expert_type = zero_expert_type + self.routed_scaling_factor = routed_scaling_factor + self.hidden_act = "silu" diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py new file mode 100644 index 000000000..0dbe37aa0 --- /dev/null +++ b/python/sglang/srt/configs/model_config.py @@ -0,0 +1,811 @@ +# Copyright 2023-2024 SGLang 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 json +import logging +import math +import os +from enum import Enum, IntEnum, auto +from typing import List, Optional, Set, Union + +import torch +from transformers import PretrainedConfig + +from sglang.srt.hf_transformers_utils import ( + get_config, + get_context_length, + get_generation_config, + get_hf_text_config, + get_sparse_attention_config, +) +from sglang.srt.layers.quantization import QUANTIZATION_METHODS +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import get_bool_env_var, is_hip +from sglang.utils import is_in_ci + +logger = logging.getLogger(__name__) + + +class AttentionArch(IntEnum): + MLA = auto() + MHA = auto() + + +class ModelImpl(str, Enum): + AUTO = "auto" + SGLANG = "sglang" + TRANSFORMERS = "transformers" + + +class ModelConfig: + def __init__( + self, + model_path: str, + trust_remote_code: bool = True, + revision: Optional[str] = None, + context_length: Optional[int] = None, + model_override_args: str = "{}", + is_embedding: Optional[bool] = None, + enable_multimodal: Optional[bool] = None, + dtype: str = "auto", + quantization: Optional[str] = None, + override_config_file: Optional[str] = None, + is_draft_model: bool = False, + hybrid_kvcache_ratio: Optional[float] = None, + model_impl: Union[str, ModelImpl] = ModelImpl.AUTO, + ) -> None: + # Parse args + self.model_path = model_path + self.revision = revision + self.quantization = quantization + self.model_impl = model_impl + + self.maybe_pull_model_tokenizer_from_remote() + self.model_override_args = json.loads(model_override_args) + kwargs = {} + if override_config_file and override_config_file.strip(): + kwargs["_configuration_file"] = override_config_file.strip() + + self.hf_config = get_config( + self.model_path, + trust_remote_code=trust_remote_code, + revision=revision, + model_override_args=self.model_override_args, + **kwargs, + ) + + self.hf_generation_config = get_generation_config( + self.model_path, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + + self.hf_text_config = get_hf_text_config(self.hf_config) + self.attention_chunk_size = getattr( + self.hf_text_config, "attention_chunk_size", None + ) + self.is_hybrid = is_hybrid_model( + self.hf_config.architectures, + hybrid_kvcache_ratio=hybrid_kvcache_ratio, + context_length=context_length, + attention_chunk_size=self.attention_chunk_size, + ) + if self.is_hybrid is not None: + self.swa_attention_layer_ids, self.full_attention_layer_ids = ( + get_hybrid_layer_ids( + self.hf_config.architectures, self.hf_text_config.num_hidden_layers + ) + ) + + if enable_multimodal is None: + mm_disabled_models = [ + "Gemma3ForConditionalGeneration", + "Llama4ForConditionalGeneration", + "Step3VLForConditionalGeneration", + ] + if self.hf_config.architectures[0] in mm_disabled_models: + enable_multimodal = False + logger.info( + f"Multimodal is disabled for {self.hf_config.model_type}. To enable it, set --enable-multimodal." + ) + else: + enable_multimodal = True + + if ( + is_draft_model + and self.hf_config.architectures[0] == "DeepseekV3ForCausalLM" + ): + self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" + + if is_draft_model and self.hf_config.architectures[0] == "Glm4MoeForCausalLM": + self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN" + + if ( + is_draft_model + and self.hf_config.architectures[0] == "LongcatFlashForCausalLM" + ): + self.hf_config.architectures[0] = "LongcatFlashForCausalLMNextN" + self.hf_config.num_hidden_layers = self.hf_config.num_nextn_predict_layers + + if is_draft_model and self.hf_config.architectures[0] == "MiMoForCausalLM": + self.hf_config.architectures[0] = "MiMoMTP" + if ( + is_draft_model + and self.hf_config.architectures[0] == "Ernie4_5_MoeForCausalLM" + ): + self.hf_config.architectures[0] = "Ernie4_5_MoeForCausalLMMTP" + + if is_draft_model and self.hf_config.architectures[0] == "Qwen3NextForCausalLM": + self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP" + + # Check model type + self.is_generation = is_generation_model( + self.hf_config.architectures, is_embedding + ) + self.is_multimodal = enable_multimodal and is_multimodal_model( + self.hf_config.architectures + ) + self.is_multimodal_gen = enable_multimodal and is_multimodal_gen_model( + self.hf_config.architectures + ) + self.is_image_gen = enable_multimodal and is_image_gen_model( + self.hf_config.architectures + ) + self.is_audio_model = enable_multimodal and is_audio_model( + self.hf_config.architectures + ) + self.is_multimodal_chunked_prefill_supported = ( + enable_multimodal + and is_multimodal_chunked_prefill_supported(self.hf_config.architectures) + ) + self.is_encoder_decoder = is_encoder_decoder_model(self.hf_config.architectures) + self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype) + + # Derive context length + derived_context_len = get_context_length(self.hf_text_config) + if context_length is not None: + if context_length > derived_context_len: + reason = "Target model's" if is_draft_model else "User-specified" + msg = ( + f"Warning: {reason} context_length ({context_length}) is greater than the derived context_length ({derived_context_len}). " + f"This may lead to incorrect model outputs or CUDA errors. Note that the derived context_length may differ from max_position_embeddings in the model's config." + ) + if ( + get_bool_env_var("SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN") + or is_in_ci() # FIXME: fix this special case + ): + logger.warning(msg) + self.context_len = context_length + else: + raise ValueError( + f"{msg} To allow overriding this maximum, set the env var SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1" + ) + else: + self.context_len = context_length + else: + self.context_len = derived_context_len + + # Unify the config keys for hf_text_config + self.head_dim = getattr( + self.hf_text_config, + "head_dim", + self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads, + ) + + # FIXME: temporary special judge for MLA architecture + if ( + "DeepseekV2ForCausalLM" in self.hf_config.architectures + or "DeepseekV3ForCausalLM" in self.hf_config.architectures + or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures + or "LongcatFlashForCausalLM" in self.hf_config.architectures + or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures + ): + self.head_dim = 256 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_config.kv_lora_rank + self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim + self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim + self.v_head_dim = self.hf_config.v_head_dim + + # Handle rope scaling with yarn + self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) + if self.hf_config.rope_scaling: + mscale_all_dim = self.hf_config.rope_scaling.get( + "mscale_all_dim", False + ) + scaling_factor = self.hf_config.rope_scaling["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + + elif "MiniCPM3ForCausalLM" in self.hf_config.architectures: + self.head_dim = 128 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_config.kv_lora_rank + self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim + elif "DeepseekVL2ForCausalLM" in self.hf_config.architectures and getattr( + self.hf_text_config, "use_mla", True + ): + self.head_dim = 256 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_text_config.kv_lora_rank + self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim + elif "KimiVLForConditionalGeneration" in self.hf_config.architectures: + self.head_dim = 256 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_text_config.kv_lora_rank + self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim + self.v_head_dim = self.hf_text_config.v_head_dim + self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim + else: + if ( + "MistralModel" in self.hf_config.architectures + or "MixtralForCausalLM" in self.hf_config.architectures + or "MistralForCausalLM" in self.hf_config.architectures + ): + if getattr(self, "head_dim", None) is None: + self.head_dim = ( + self.hf_config.hidden_size // self.hf_config.num_attention_heads + ) + # In transformers==4.52.3, the head_dim is null in MistralConfig + if ( + not hasattr(self.hf_text_config, "head_dim") + or self.hf_text_config.head_dim is None + ): + setattr(self.hf_text_config, "head_dim", self.head_dim) + + self.attention_arch = AttentionArch.MHA + + self.num_attention_heads = self.hf_text_config.num_attention_heads + self.num_key_value_heads = getattr( + self.hf_text_config, "num_key_value_heads", None + ) + + # for Dbrx and MPT models + if self.hf_config.model_type in ["dbrx", "mpt"]: + self.num_key_value_heads = getattr( + self.hf_config.attn_config, "kv_n_heads", None + ) + + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + self.hidden_size = self.hf_text_config.hidden_size + self.num_hidden_layers = self.hf_text_config.num_hidden_layers + self.num_attention_layers = self.num_hidden_layers + if "LongcatFlashForCausalLM" in self.hf_config.architectures: + self.num_attention_layers = self.num_hidden_layers * 2 + self.num_nextn_predict_layers = getattr( + self.hf_text_config, "num_nextn_predict_layers", None + ) + self.vocab_size = self.hf_text_config.vocab_size + + # Verify quantization + self._verify_quantization() + + # Verify dual-chunk attention config + self._verify_dual_chunk_attention_config() + + # Cache attributes + self.hf_eos_token_id = self.get_hf_eos_token_id() + + # multimodal + self.image_token_id = getattr( + self.hf_config, "image_token_id", None + ) or getattr(self.hf_config, "image_token_index", None) + + @staticmethod + def from_server_args( + server_args: ServerArgs, + model_path: str = None, + model_revision: str = None, + **kwargs, + ): + return ModelConfig( + model_path=model_path or server_args.model_path, + trust_remote_code=server_args.trust_remote_code, + revision=model_revision or server_args.revision, + context_length=server_args.context_length, + model_override_args=server_args.json_model_override_args, + is_embedding=server_args.is_embedding, + enable_multimodal=server_args.enable_multimodal, + dtype=server_args.dtype, + quantization=server_args.quantization, + hybrid_kvcache_ratio=server_args.hybrid_kvcache_ratio, + model_impl=server_args.model_impl, + **kwargs, + ) + + def get_total_num_attention_heads(self) -> int: + return self.num_attention_heads + + def get_num_attention_heads(self, tensor_parallel_size) -> int: + total_num_attention_heads = self.num_attention_heads + return max(1, total_num_attention_heads // tensor_parallel_size) + + # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L289 + def get_total_num_kv_heads(self) -> int: + """Returns the total number of KV heads.""" + # For GPTBigCode & Falcon: + # NOTE: for falcon, when new_decoder_architecture is True, the + # multi_query flag is ignored and we use n_head_kv for the number of + # KV heads. + falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"] + new_decoder_arch_falcon = ( + self.hf_config.model_type in falcon_model_types + and getattr(self.hf_config, "new_decoder_architecture", False) + ) + if not new_decoder_arch_falcon and getattr( + self.hf_text_config, "multi_query", False + ): + # Multi-query attention, only one KV head. + # Currently, tensor parallelism is not supported in this case. + return 1 + + # For DBRX and MPT + if self.hf_config.model_type in ["mpt"]: + if "kv_n_heads" in self.hf_config.attn_config: + return self.hf_config.attn_config["kv_n_heads"] + return self.hf_config.num_attention_heads + if self.hf_config.model_type in ["dbrx"]: + return getattr( + self.hf_config.attn_config, + "kv_n_heads", + self.hf_config.num_attention_heads, + ) + if self.hf_config.model_type in ["nemotron-nas"]: + nkvh = { + self.hf_config.num_attention_heads // block.attention.n_heads_in_group + for block in self.hf_config.block_configs + if not block.attention.no_op + } + if len(nkvh) == 0: + raise RuntimeError("Couldn't determine number of kv heads") + if len(nkvh) > 1: + raise ValueError( + "Variable GQA (VGQA) is not yet supported for nemotron-nas in sglang" + ) + return next(iter(nkvh)) + + attributes = [ + # For Falcon: + "n_head_kv", + "num_kv_heads", + # For LLaMA-2: + "num_key_value_heads", + # For ChatGLM: + "multi_query_group_num", + # For Step3 + "num_attention_groups", + ] + for attr in attributes: + num_kv_heads = getattr(self.hf_text_config, attr, None) + if num_kv_heads is not None: + return num_kv_heads + + # For non-grouped-query attention models, the number of KV heads is + # equal to the number of attention heads. + return self.hf_text_config.num_attention_heads + + def get_num_kv_heads(self, tensor_parallel_size) -> int: + """Returns the number of KV heads per GPU.""" + total_num_kv_heads = self.get_total_num_kv_heads() + # If tensor parallelism is used, we divide the number of KV heads by + # the tensor parallel size. We will replicate the KV heads in the + # case where the number of KV heads is smaller than the tensor + # parallel size so each GPU has at least one KV head. + return max(1, total_num_kv_heads // tensor_parallel_size) + + # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py + def _parse_quant_hf_config(self): + quant_cfg = getattr(self.hf_config, "quantization_config", None) + if quant_cfg is None: + # compressed-tensors uses a "compression_config" key + quant_cfg = getattr(self.hf_config, "compression_config", None) + if quant_cfg is None: + # check if is modelopt or mixed-precision model -- Both of them don't have corresponding field + # in hf `config.json` but has a standalone `hf_quant_config.json` in the root directory + # example: https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8/tree/main + # example: https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/tree/main + is_local = os.path.exists(self.model_path) + modelopt_quant_config = {"quant_method": "modelopt"} + if not is_local: + import huggingface_hub + + try: + from huggingface_hub import HfApi + + hf_api = HfApi() + if hf_api.file_exists(self.model_path, "hf_quant_config.json"): + quant_cfg = modelopt_quant_config + except huggingface_hub.errors.OfflineModeIsEnabled: + logger.warning( + "Offline mode is enabled, skipping hf_quant_config.json check" + ) + pass + + elif os.path.exists(os.path.join(self.model_path, "hf_quant_config.json")): + quant_config_file = os.path.join( + self.model_path, "hf_quant_config.json" + ) + with open(quant_config_file) as f: + quant_config_dict = json.load(f) + json_quant_configs = quant_config_dict["quantization"] + quant_algo = json_quant_configs.get("quant_algo", None) + if quant_algo == "MIXED_PRECISION": + quant_cfg = {"quant_method": "w4afp8"} + else: + quant_cfg = modelopt_quant_config + return quant_cfg + + # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py + def _verify_quantization(self) -> None: + supported_quantization = [*QUANTIZATION_METHODS] + rocm_supported_quantization = [ + "awq", + "gptq", + "fp8", + "compressed_tensors", + "compressed-tensors", + "fbgemm_fp8", + "w8a8_fp8", + "petit_nvfp4", + "quark", + "mxfp4", + ] + optimized_quantization_methods = [ + "fp8", + "marlin", + "modelopt", + "gptq_marlin_24", + "gptq_marlin", + "awq_marlin", + "fbgemm_fp8", + "compressed_tensors", + "compressed-tensors", + "experts_int8", + "w8a8_int8", + "w8a8_fp8", + "moe_wna16", + "qoq", + "w4afp8", + "petit_nvfp4", + ] + compatible_quantization_methods = { + "modelopt_fp4": ["modelopt"], + "petit_nvfp4": ["modelopt"], + "w8a8_int8": ["compressed-tensors", "compressed_tensors"], + "w8a8_fp8": ["compressed-tensors", "compressed_tensors"], + } + if self.quantization is not None: + self.quantization = self.quantization.lower() + + # Parse quantization method from the HF model config, if available. + quant_cfg = self._parse_quant_hf_config() + + if quant_cfg is not None: + quant_method = quant_cfg.get( + "quant_method", "" if not self.quantization else self.quantization + ).lower() + + # Detect which checkpoint is it + for _, method in QUANTIZATION_METHODS.items(): + quantization_override = method.override_quantization_method( + quant_cfg, self.quantization + ) + if quantization_override: + quant_method = quantization_override + self.quantization = quantization_override + break + + # Verify quantization configurations. + if self.quantization is None: + self.quantization = quant_method + elif self.quantization != quant_method: + if ( + self.quantization not in compatible_quantization_methods + or quant_method + not in compatible_quantization_methods[self.quantization] + ): + raise ValueError( + "Quantization method specified in the model config " + f"({quant_method}) does not match the quantization " + f"method specified in the `quantization` argument " + f"({self.quantization})." + ) + + if self.quantization is not None: + if self.quantization not in supported_quantization: + raise ValueError( + f"Unknown quantization method: {self.quantization}. Must " + f"be one of {supported_quantization}." + ) + if is_hip() and self.quantization not in rocm_supported_quantization: + raise ValueError( + f"{self.quantization} quantization is currently not " + f"supported in ROCm." + ) + if self.quantization not in optimized_quantization_methods: + logger.warning( + "%s quantization is not fully " + "optimized yet. The speed can be slower than " + "non-quantized models.", + self.quantization, + ) + + def _verify_dual_chunk_attention_config(self) -> None: + if hasattr(self.hf_config, "dual_chunk_attention_config"): + # Try loading the sparse attention config + sparse_attn_config = get_sparse_attention_config(self.model_path) + if not sparse_attn_config: + return + self.hf_config.dual_chunk_attention_config["sparse_attention_config"] = ( + sparse_attn_config + ) + if ( + "sparse_attention_enabled" + not in self.hf_config.dual_chunk_attention_config + ): + self.hf_config.dual_chunk_attention_config[ + "sparse_attention_enabled" + ] = True + + def get_hf_eos_token_id(self) -> Optional[Set[int]]: + eos_ids = getattr(self.hf_config, "eos_token_id", None) + if eos_ids is not None: + # it can be either int or list of int + eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids) + if eos_ids is None: + eos_ids = set() + if self.hf_generation_config: + generation_eos_ids = getattr( + self.hf_generation_config, "eos_token_id", None + ) + if generation_eos_ids: + generation_eos_ids = ( + {generation_eos_ids} + if isinstance(generation_eos_ids, int) + else set(generation_eos_ids) + ) + eos_ids = eos_ids | generation_eos_ids + return eos_ids + + def maybe_pull_model_tokenizer_from_remote(self) -> None: + """ + Pull the model config files to a temporary + directory in case of remote. + + Args: + model: The model name or path. + + """ + from sglang.srt.connector import create_remote_connector + from sglang.srt.utils import is_remote_url + + if is_remote_url(self.model_path): + logger.info("Pulling model configs from remote...") + # BaseConnector implements __del__() to clean up the local dir. + # Since config files need to exist all the time, so we DO NOT use + # with statement to avoid closing the client. + client = create_remote_connector(self.model_path) + if is_remote_url(self.model_path): + client.pull_files(allow_pattern=["*config.json"]) + self.model_weights = self.model_path + self.model_path = client.get_local_dir() + + +# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py +_STR_DTYPE_TO_TORCH_DTYPE = { + "half": torch.float16, + "float16": torch.float16, + "float": torch.float32, + "float32": torch.float32, + "bfloat16": torch.bfloat16, +} + + +# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py +def _get_and_verify_dtype( + config: PretrainedConfig, + dtype: Union[str, torch.dtype], +) -> torch.dtype: + # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct + # because config.torch_dtype can be None. + config_dtype = getattr(config, "torch_dtype", None) + if isinstance(config_dtype, str): + config_dtype = _STR_DTYPE_TO_TORCH_DTYPE.get(config_dtype, None) + if config_dtype is None: + config_dtype = torch.float32 + + if isinstance(dtype, str): + dtype = dtype.lower() + if dtype == "auto": + if config_dtype == torch.float32: + if config.model_type.startswith("gemma"): + if config.model_type == "gemma": + gemma_version = "" + else: + gemma_version = config.model_type[5] + logger.info( + f"For Gemma {gemma_version}, we downcast float32 to bfloat16 instead " + "of float16 by default. Please specify `dtype` if you " + "want to use float16." + ) + torch_dtype = torch.bfloat16 + else: + # Following the common practice, we use float16 for float32 + # models. + torch_dtype = torch.float16 + else: + torch_dtype = config_dtype + else: + if dtype not in _STR_DTYPE_TO_TORCH_DTYPE: + raise ValueError(f"Unknown dtype: {dtype}") + torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype] + elif isinstance(dtype, torch.dtype): + torch_dtype = dtype + else: + raise ValueError(f"Unknown dtype: {dtype}") + + # Verify the dtype. + if torch_dtype != config_dtype: + if torch_dtype == torch.float32: + # Upcasting to float32 is allowed. + logger.info("Upcasting %s to %s.", config_dtype, torch_dtype) + pass + elif config_dtype == torch.float32: + # Downcasting from float32 to float16 or bfloat16 is allowed. + logger.info("Downcasting %s to %s.", config_dtype, torch_dtype) + pass + else: + # Casting between float16 and bfloat16 is allowed with a warning. + logger.warning("Casting %s to %s.", config_dtype, torch_dtype) + + return torch_dtype + + +def is_generation_model(model_architectures: List[str], is_embedding: bool = False): + # We have two ways to determine whether a model is a generative model. + # 1. Check the model architecture + # 2. check the `is_embedding` server args + + if ( + "LlamaEmbeddingModel" in model_architectures + or "MistralModel" in model_architectures + or "LlamaForSequenceClassification" in model_architectures + or "LlamaForSequenceClassificationWithNormal_Weights" in model_architectures + or "InternLM2ForRewardModel" in model_architectures + or "Qwen2ForRewardModel" in model_architectures + or "Qwen2ForSequenceClassification" in model_architectures + or "Qwen3ForSequenceClassification" in model_architectures + or "CLIPModel" in model_architectures + or "BertModel" in model_architectures + or "Contriever" in model_architectures + or "BertForSequenceClassification" in model_architectures + or "XLMRobertaModel" in model_architectures + or "XLMRobertaForSequenceClassification" in model_architectures + ): + return False + else: + return not is_embedding + + +multimodal_model_archs = [ + "CLIPModel", + "DeepseekVL2ForCausalLM", + "Gemma3ForConditionalGeneration", + "Gemma3nForConditionalGeneration", + "Glm4vForConditionalGeneration", + "Glm4vMoeForConditionalGeneration", + "Grok1VForCausalLM", + "Grok1AForCausalLM", + "LlavaLlamaForCausalLM", + "Llama4ForConditionalGeneration", + "LlavaMistralForCausalLM", + "LlavaQwenForCausalLM", + "LlavaForConditionalGeneration", + "LlavaVidForCausalLM", + "MiniCPMO", + "MiniCPMV", + "Mistral3ForConditionalGeneration", + "MultiModalityCausalLM", + "MllamaForConditionalGeneration", + "Qwen2AudioForConditionalGeneration", + "Qwen2VLForConditionalGeneration", + "Qwen2_5_VLForConditionalGeneration", + "KimiVLForConditionalGeneration", + "InternVLChatModel", + "InternS1ForConditionalGeneration", + "Phi4MMForCausalLM", + "VILAForConditionalGeneration", + "Step3VLForConditionalGeneration", +] + + +def is_multimodal_model(model_architectures: List[str]): + if any( + multi_model_arch in model_architectures + for multi_model_arch in multimodal_model_archs + ): + return True + else: + return False + + +def is_multimodal_gen_model(model_architectures: List[str]): + return False + + +def is_image_gen_model(model_architectures: List[str]): + return False + + +def is_audio_model(model_architectures: List[str]): + return False + + +def is_encoder_decoder_model(model_architectures: List[str]): + return "MllamaForConditionalGeneration" in model_architectures + + +def is_multimodal_chunked_prefill_supported(model_architectures: List[str]): + """Check if chunked prefill is supported for a MultiModal model.""" + unsupported = [ + "Grok1VForCausalLM", + "Grok1AForCausalLM", + "LlavaLlamaForCausalLM", + "MllamaForConditionalGeneration", + "CLIPModel", + ] + if any(multi_model_arch in unsupported for multi_model_arch in model_architectures): + return False + else: + return True + + +def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def is_hybrid_model( + model_architectures: List[str], + hybrid_kvcache_ratio: Optional[float], + context_length: Optional[int], + attention_chunk_size: Optional[int], +): + if hybrid_kvcache_ratio is None: + return None + elif ( + hybrid_kvcache_ratio > 0 + and model_architectures[0] == "Llama4ForConditionalGeneration" + and context_length > attention_chunk_size + ): + return hybrid_kvcache_ratio + else: + return None + + +def get_hybrid_layer_ids(model_architectures: List[str], num_hidden_layers: int): + if "Llama4ForConditionalGeneration" in model_architectures: + swa_attention_layer_ids = [ + i for i in range(num_hidden_layers) if (i + 1) % 4 != 0 + ] + full_attention_layer_ids = [ + i for i in range(num_hidden_layers) if (i + 1) % 4 == 0 + ] + else: + swa_attention_layer_ids = None + full_attention_layer_ids = None + return swa_attention_layer_ids, full_attention_layer_ids diff --git a/python/sglang/srt/configs/qwen3_next.py b/python/sglang/srt/configs/qwen3_next.py new file mode 100644 index 000000000..099d14d41 --- /dev/null +++ b/python/sglang/srt/configs/qwen3_next.py @@ -0,0 +1,326 @@ +# coding=utf-8 +# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. 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. +"""Qwen3Hybrid model configuration""" + +import enum +import os + +import numpy as np +import torch +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_rope_utils import rope_config_validation +from transformers.utils import logging + +from sglang.srt.distributed.utils import divide +from sglang.srt.layers.dp_attention import get_attention_tp_size + +logger = logging.get_logger(__name__) + + +# NOTE: HybridLayerType +class HybridLayerType(enum.Enum): + full_attention = "attention" + swa_attention = "swa_attention" + linear_attention = "linear_attention" + mamba2 = "mamba" + + +class Qwen3NextConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`Qwen3NextModel`]. It is used to instantiate a + Qwen3-Next model according to the specified arguments, defining the model architecture. + Instantiating a configuration with the defaults will yield a similar configuration to that of + Qwen3-Next-80B-A3B-Instruct [Qwen/Qwen3-Next-80B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct). + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 151936): + Vocabulary size of the model. Defines the number of different tokens that can be represented by the + `inputs_ids`. + hidden_size (`int`, *optional*, defaults to 2048): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 5632): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 48): + Number of hidden layers in the Transformer encoder. + num_attention_heads (`int`, *optional*, defaults to 16): + Number of attention heads for each attention layer in the Transformer encoder. + num_key_value_heads (`int`, *optional*, defaults to 2): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`. + hidden_act (`str`, *optional*, defaults to `"silu"`): + The non-linear activation function in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 32768): + The maximum sequence length that this model might ever be used with. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether the model's input and output word embeddings should be tied. + rope_theta (`float`, *optional*, defaults to 10000.0): + The base period of the RoPE embeddings. + rope_scaling (`Dict`, *optional*): + Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type + and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value + accordingly. + Expected contents: + `rope_type` (`str`): + The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', + 'llama3'], with 'default' being the original RoPE implementation. + `factor` (`float`, *optional*): + Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In + most scaling types, a `factor` of x will enable the model to handle sequences of length x * + original maximum pre-trained length. + `original_max_position_embeddings` (`int`, *optional*): + Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during + pretraining. + `attention_factor` (`float`, *optional*): + Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention + computation. If unspecified, it defaults to value recommended by the implementation, using the + `factor` field to infer the suggested value. + `beta_fast` (`float`, *optional*): + Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear + ramp function. If unspecified, it defaults to 32. + `beta_slow` (`float`, *optional*): + Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear + ramp function. If unspecified, it defaults to 1. + `short_factor` (`List[float]`, *optional*): + Only used with 'longrope'. The scaling factor to be applied to short contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `long_factor` (`List[float]`, *optional*): + Only used with 'longrope'. The scaling factor to be applied to long contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `low_freq_factor` (`float`, *optional*): + Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE + `high_freq_factor` (`float`, *optional*): + Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE + partial_rotary_factor (`float`, *optional*, defaults to 0.25): + Percentage of the query and keys which will have rotary embedding. + attention_bias (`bool`, *optional*, defaults to `False`): + Whether to use a bias in the query, key, value and output projection layers during self-attention. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + head_dim (`int`, *optional*, defaults to 256): + Projection weights dimension in multi-head attention. + linear_conv_kernel_dim (`int`, *optional*, defaults to 4): + Kernel size of the convolution used in linear attention layers. + linear_key_head_dim (`int`, *optional*, defaults to 128): + Dimension of each key head in linear attention. + linear_value_head_dim (`int`, *optional*, defaults to 128): + Dimension of each value head in linear attention. + linear_num_key_heads (`int`, *optional*, defaults to 16): + Number of key heads used in linear attention layers. + linear_num_value_heads (`int`, *optional*, defaults to 32): + Number of value heads used in linear attention layers. + decoder_sparse_step (`int`, *optional*, defaults to 1): + The frequency of the MoE layer. + moe_intermediate_size (`int`, *optional*, defaults to 512): + Intermediate size of the routed expert. + shared_expert_intermediate_size (`int`, *optional*, defaults to 512): + Intermediate size of the shared expert. + num_experts_per_tok (`int`, *optional*, defaults to 10): + Number of selected experts. + num_experts (`int`, *optional*, defaults to 512): + Number of routed experts. + norm_topk_prob (`bool`, *optional*, defaults to `True`): + Whether to normalize the topk probabilities. + output_router_logits (`bool`, *optional*, defaults to `False`): + Whether or not the router logits should be returned by the model. Enabling this will also + allow the model to output the auxiliary loss, including load balancing loss and router z-loss. + router_aux_loss_coef (`float`, *optional*, defaults to 0.001): + The aux loss factor for the total loss. + mlp_only_layers (`list[int]`, *optional*, defaults to `[]`): + Indicate which layers use Qwen3NextMLP rather than Qwen3NextSparseMoeBlock + The list contains layer index, from 0 to num_layers-1 if we have num_layers layers + If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity. + layer_types (`list[str]`, *optional*, defaults to None): + Types of each layer (attention or linear). + + ```python + >>> from transformers import Qwen3NextModel, Qwen3NextConfig + + >>> # Initializing a Qwen3Next style configuration + >>> configuration = Qwen3NextConfig() + + >>> # Initializing a model from the Qwen3-Next-80B-A3B style configuration + >>> model = Qwen3NextModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "qwen3_next" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=151936, + hidden_size=2048, + intermediate_size=5632, + num_hidden_layers=48, + num_attention_heads=16, + num_key_value_heads=2, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + partial_rotary_factor=0.25, + attention_bias=False, + attention_dropout=0.0, + head_dim=256, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + decoder_sparse_step=1, + moe_intermediate_size=512, + shared_expert_intermediate_size=512, + num_experts_per_tok=10, + num_experts=512, + norm_topk_prob=True, + output_router_logits=False, + router_aux_loss_coef=0.001, + mlp_only_layers=[], + layer_types=None, + **kwargs, + ): + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.partial_rotary_factor = partial_rotary_factor + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.head_dim = head_dim + rope_config_validation(self) + + # linear attention (gdn now part) + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_key_head_dim = linear_key_head_dim + self.linear_value_head_dim = linear_value_head_dim + self.linear_num_key_heads = linear_num_key_heads + self.linear_num_value_heads = linear_num_value_heads + + # MoE arguments + self.decoder_sparse_step = decoder_sparse_step + self.moe_intermediate_size = moe_intermediate_size + self.shared_expert_intermediate_size = shared_expert_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.norm_topk_prob = norm_topk_prob + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + self.mlp_only_layers = mlp_only_layers + + @property + def layers_block_type(self): + layer_type_list = [] + + for l in range(self.num_hidden_layers): + if (l + 1) % self.full_attention_interval == 0: + layer_type_list.append(HybridLayerType.full_attention.value) + else: + layer_type_list.append(HybridLayerType.linear_attention.value) + + return layer_type_list + + @property + def linear_layer_ids(self): + return [ + i + for i, type_value in enumerate(self.layers_block_type) + if type_value == HybridLayerType.linear_attention.value + ] + + @property + def full_attention_layer_ids(self): + return [ + i + for i, type_value in enumerate(self.layers_block_type) + if type_value == HybridLayerType.full_attention.value + ] + + @property + def hybrid_gdn_params(self): + world_size = get_attention_tp_size() + conv_dim = ( + self.linear_key_head_dim * self.linear_num_key_heads * 2 + + self.linear_value_head_dim * self.linear_num_value_heads + ) + conv_state_shape = ( + divide(conv_dim, world_size), + self.linear_conv_kernel_dim - 1, + ) + + temporal_state_shape = ( + divide(self.linear_num_value_heads, world_size), + self.linear_key_head_dim, + self.linear_value_head_dim, + ) + conv_dtype = torch.bfloat16 + dtype_map = { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + } + ssm_dtype = dtype_map[os.environ["SGLANG_MAMBA_SSM_DTYPE"]] + mamba_layers = self.linear_layer_ids + return ( + conv_state_shape, + temporal_state_shape, + conv_dtype, + ssm_dtype, + mamba_layers, + ) + + @property + def mamba_cache_per_req(self): + conv_state_shape, temporal_state_shape, conv_dtype, ssm_dtype, mamba_layers = ( + self.hybrid_gdn_params + ) + mamba_layers_len = len(mamba_layers) + + return ( + int(np.prod(conv_state_shape)) * conv_dtype.itemsize + + int(np.prod(temporal_state_shape)) * ssm_dtype.itemsize + ) * mamba_layers_len diff --git a/python/sglang/srt/configs/step3_vl.py b/python/sglang/srt/configs/step3_vl.py new file mode 100644 index 000000000..5519605c6 --- /dev/null +++ b/python/sglang/srt/configs/step3_vl.py @@ -0,0 +1,172 @@ +from typing import Any, Optional, Union + +from transformers.configuration_utils import PretrainedConfig + + +class Step3VisionEncoderConfig(PretrainedConfig): + model_type = "step3_vision_encoder" + + def __init__( + self, + hidden_size=1792, + intermediate_size=3072, + output_hidden_size=4096, + num_hidden_layers=63, + num_attention_heads=16, + num_channels=3, + image_size=728, + patch_size=14, + hidden_act="quick_gelu", + layer_norm_eps=1e-5, + **kwargs, + ): + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.output_hidden_size = output_hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_channels = num_channels + self.patch_size = patch_size + self.image_size = image_size + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + super().__init__(**kwargs) + + +class Step3TextConfig(PretrainedConfig): + model_type = "step3_text" + architectures = ["Step3TextForCausalLM"] + + def __init__( + self, + hidden_size: int = 7168, + intermediate_size: int = 18432, + num_attention_heads: int = 64, + num_attention_groups: int = 1, + num_hidden_layers: int = 61, + max_seq_len: int = 65536, + vocab_size: int = 128815, + rms_norm_eps: float = 1e-5, + moe_intermediate_size: int = 5120, + moe_num_experts: int = 48, + moe_top_k: int = 3, + rope_theta: float = 500000, + rope_scaling: Optional[dict[str, Any]] = None, + max_position_embedding: int = 65536, + share_expert_dim: int = 5120, + share_q_dim: int = 2048, + head_dim: int = 256, + norm_expert_weight: bool = False, + moe_layers_enum: tuple[int] = ( + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 40, + 41, + 42, + 43, + 44, + 45, + 46, + 47, + 48, + 49, + 50, + 51, + 52, + 53, + 54, + 55, + 56, + 57, + 58, + 59, + ), + **kwargs, + ) -> None: + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_attention_heads = num_attention_heads + self.num_attention_groups = num_attention_groups + self.num_hidden_layers = num_hidden_layers + self.max_seq_len = max_seq_len + self.vocab_size = vocab_size + self.rms_norm_eps = rms_norm_eps + self.moe_intermediate_size = moe_intermediate_size + self.moe_num_experts = moe_num_experts + self.moe_top_k = moe_top_k + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.max_position_embedding = max_position_embedding + self.share_expert_dim = share_expert_dim + self.share_q_dim = share_q_dim + self.head_dim = head_dim + self.norm_expert_weight = norm_expert_weight + self.moe_layers_enum = moe_layers_enum + + super().__init__(**kwargs) + + +class Step3VLConfig(PretrainedConfig): + model_type = "step3_vl" + + def __init__( + self, + vision_config: Optional[Union[dict, Step3VisionEncoderConfig]] = None, + text_config: Optional[Union[dict, Step3TextConfig]] = None, + understand_projector_stride: int = 1, + projector_bias: bool = True, + image_token_id: int = 128001, + **kwargs, + ) -> None: + if vision_config is None: + vision_config = Step3VisionEncoderConfig() + elif isinstance(vision_config, dict): + vision_config = Step3VisionEncoderConfig(**vision_config) + self.vision_config = vision_config + + if text_config is None: + text_config = Step3TextConfig() + elif isinstance(text_config, dict): + text_config = Step3TextConfig(**text_config) + self.text_config = text_config + + self.understand_projector_stride = understand_projector_stride + self.projector_bias = projector_bias + self.hidden_size = text_config.hidden_size + self.image_token_id = image_token_id + + super().__init__(**kwargs) diff --git a/python/sglang/srt/configs/update_config.py b/python/sglang/srt/configs/update_config.py new file mode 100644 index 000000000..abbd724fb --- /dev/null +++ b/python/sglang/srt/configs/update_config.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +DEFAULT_MOE_PADDING_SIZE = 32 + + +if TYPE_CHECKING: + from sglang.srt.configs.load_config import LoadConfig + from sglang.srt.configs.model_config import ModelConfig + + +def may_get_weight_block_size(model_config, load_config): + from sglang.srt.model_loader.loader import _get_quantization_config + from sglang.srt.model_loader.utils import get_model_architecture + + model_class, _ = get_model_architecture(model_config) + packed_modules_mapping = getattr(model_class, "packed_modules_mapping", {}) + + quant_config = _get_quantization_config( + model_config, load_config, packed_modules_mapping + ) + + if quant_config is not None and hasattr(quant_config, "weight_block_size"): + return getattr(quant_config, "weight_block_size") + return None + + +def get_moe_padding_size(weight_block_size): + if weight_block_size is not None: + # See NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n. + assert ( + len(weight_block_size) == 2 + ), "Only len(weight_block_size) == 2 is supported" + assert ( + weight_block_size[0] == weight_block_size[1] + ), "Only weight_block_size[0] == weight_block_size[1] is supported" + + return weight_block_size[0] + + return DEFAULT_MOE_PADDING_SIZE + + +def get_num_heads_padding_size(tp_size, weight_block_size): + pad_size = ( + tp_size * 2 if tp_size % 2 == 1 and weight_block_size is not None else tp_size + ) + return pad_size + + +def update_intermediate_size(model_config, attr_name, intermediate_padding_size): + attr_value = intermediate_padding_size + if hasattr(model_config, "hf_config") and hasattr( + model_config.hf_config, attr_name + ): + attr_value = getattr(model_config.hf_config, attr_name) + elif hasattr(model_config, attr_name): + attr_value = getattr(model_config, attr_name) + + if attr_value % intermediate_padding_size != 0: + from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size + + attr_value = pad_vocab_size(attr_value, intermediate_padding_size) + if hasattr(model_config, "hf_config"): + setattr(model_config.hf_config, attr_name, attr_value) + if hasattr(model_config, "hf_text_config"): + setattr(model_config.hf_text_config, attr_name, attr_value) + else: + setattr(model_config, attr_name, attr_value) + + return model_config + + +def adjust_config_with_unaligned_cpu_tp( + model_config: ModelConfig, load_config: LoadConfig, tp_size: int +) -> ModelConfig: + # Support the case where the num_attention_heads is not divisible by the TP size. + weight_block_size = may_get_weight_block_size(model_config, load_config) + + model_config.hf_config.original_num_attention_heads = ( + model_config.num_attention_heads + ) + model_config.hf_text_config.original_num_attention_heads = ( + model_config.num_attention_heads + ) + + model_config.hf_config.original_total_num_kv_heads = ( + model_config.get_total_num_kv_heads() + ) + model_config.hf_text_config.original_total_num_kv_heads = ( + model_config.get_total_num_kv_heads() + ) + + if ( + model_config.num_attention_heads % tp_size != 0 + or model_config.get_total_num_kv_heads() % tp_size != 0 + ): + # Compute the head_dim using the model_config.num_attention_heads before padding + if not hasattr(model_config.hf_config, "head_dim"): + model_config.hf_config.head_dim = ( + model_config.hidden_size // model_config.num_attention_heads + ) + + query_heads_per_kv = ( + model_config.num_attention_heads // model_config.get_total_num_kv_heads() + ) + total_kv_heads = model_config.get_total_num_kv_heads() + from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size + + pad_size = get_num_heads_padding_size(tp_size, weight_block_size) + num_key_value_heads = pad_vocab_size(total_kv_heads, pad_size) + + model_config.num_key_value_heads = num_key_value_heads + model_config.hf_config.num_key_value_heads = num_key_value_heads + model_config.hf_text_config.num_key_value_heads = num_key_value_heads + + num_attention_heads = num_key_value_heads * query_heads_per_kv + model_config.num_attention_heads = num_attention_heads + model_config.hf_config.num_attention_heads = num_attention_heads + model_config.hf_text_config.num_attention_heads = num_attention_heads + + intermediate_padding_size = tp_size * get_moe_padding_size(weight_block_size) + model_config = update_intermediate_size( + model_config, "moe_intermediate_size", intermediate_padding_size + ) + model_config = update_intermediate_size( + model_config, "intermediate_size", intermediate_padding_size + ) + model_config = update_intermediate_size( + model_config, "intermediate_size_mlp", intermediate_padding_size + ) + if ( + hasattr(model_config.hf_config, "vision_config") + and model_config.hf_config.vision_config.model_type == "siglip_vision_model" + ): + model_config.hf_config.vision_config.original_num_attention_heads = ( + model_config.num_attention_heads + ) + if model_config.hf_config.vision_config.num_attention_heads % tp_size != 0: + model_config.hf_config.vision_config.head_dim = ( + model_config.hf_config.vision_config.hidden_size + // model_config.hf_config.vision_config.num_attention_heads + ) + from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size + + pad_size = get_num_heads_padding_size(tp_size, weight_block_size) + model_config.hf_config.vision_config.num_attention_heads = pad_vocab_size( + model_config.hf_config.vision_config.num_attention_heads, pad_size + ) + model_config.hf_config.vision_config = update_intermediate_size( + model_config.hf_config.vision_config, + "intermediate_size", + intermediate_padding_size, + ) + + return model_config diff --git a/python/sglang/srt/configs/utils.py b/python/sglang/srt/configs/utils.py new file mode 100644 index 000000000..8403d95da --- /dev/null +++ b/python/sglang/srt/configs/utils.py @@ -0,0 +1,25 @@ +from typing import Type + +from transformers import ( + AutoImageProcessor, + AutoProcessor, + BaseImageProcessor, + PretrainedConfig, + ProcessorMixin, +) + + +def register_image_processor( + config: Type[PretrainedConfig], image_processor: Type[BaseImageProcessor] +): + """ + register customized hf image processor while removing hf impl + """ + AutoImageProcessor.register(config, None, image_processor, None, exist_ok=True) + + +def register_processor(config: Type[PretrainedConfig], processor: Type[ProcessorMixin]): + """ + register customized hf processor while removing hf impl + """ + AutoProcessor.register(config, processor, exist_ok=True) diff --git a/python/sglang/srt/connector/__init__.py b/python/sglang/srt/connector/__init__.py new file mode 100644 index 000000000..38e1d5eab --- /dev/null +++ b/python/sglang/srt/connector/__init__.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 + +import enum +import logging + +from sglang.srt.connector.base_connector import ( + BaseConnector, + BaseFileConnector, + BaseKVConnector, +) +from sglang.srt.connector.redis import RedisConnector +from sglang.srt.connector.s3 import S3Connector +from sglang.srt.utils import parse_connector_type + +logger = logging.getLogger(__name__) + + +class ConnectorType(str, enum.Enum): + FS = "filesystem" + KV = "KV" + + +def create_remote_connector(url, **kwargs) -> BaseConnector: + connector_type = parse_connector_type(url) + if connector_type == "redis": + return RedisConnector(url) + elif connector_type == "s3": + return S3Connector(url) + else: + raise ValueError(f"Invalid connector type: {url}") + + +def get_connector_type(client: BaseConnector) -> ConnectorType: + if isinstance(client, BaseKVConnector): + return ConnectorType.KV + if isinstance(client, BaseFileConnector): + return ConnectorType.FS + + raise ValueError(f"Invalid connector type: {client}") + + +__all__ = [ + "BaseConnector", + "BaseFileConnector", + "BaseKVConnector", + "RedisConnector", + "S3Connector", + "ConnectorType", + "create_remote_connector", + "get_connector_type", +] diff --git a/python/sglang/srt/connector/base_connector.py b/python/sglang/srt/connector/base_connector.py new file mode 100644 index 000000000..c9a1c36e2 --- /dev/null +++ b/python/sglang/srt/connector/base_connector.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +import shutil +import signal +import tempfile +from abc import ABC, abstractmethod +from typing import Generator, List, Optional, Tuple + +import torch + + +class BaseConnector(ABC): + """ + For fs connector such as s3: + :/// + + For kv connector such as redis: + ://://keys/ + ://files/ + """ + + def __init__(self, url: str): + self.url = url + self.closed = False + self.local_dir = tempfile.mkdtemp() + for sig in (signal.SIGINT, signal.SIGTERM): + existing_handler = signal.getsignal(sig) + signal.signal(sig, self._close_by_signal(existing_handler)) + + def get_local_dir(self): + return self.local_dir + + @abstractmethod + def weight_iterator( + self, rank: int = 0 + ) -> Generator[Tuple[str, torch.Tensor], None, None]: + raise NotImplementedError() + + @abstractmethod + def pull_files( + self, + allow_pattern: Optional[List[str]] = None, + ignore_pattern: Optional[List[str]] = None, + ) -> None: + raise NotImplementedError() + + def close(self): + if self.closed: + return + + self.closed = True + if os.path.exists(self.local_dir): + shutil.rmtree(self.local_dir) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + + def __del__(self): + self.close() + + def _close_by_signal(self, existing_handler=None): + + def new_handler(signum, frame): + self.close() + if existing_handler: + existing_handler(signum, frame) + + return new_handler + + +class BaseKVConnector(BaseConnector): + + @abstractmethod + def get(self, key: str) -> Optional[torch.Tensor]: + raise NotImplementedError() + + @abstractmethod + def getstr(self, key: str) -> Optional[str]: + raise NotImplementedError() + + @abstractmethod + def set(self, key: str, obj: torch.Tensor) -> None: + raise NotImplementedError() + + @abstractmethod + def setstr(self, key: str, obj: str) -> None: + raise NotImplementedError() + + @abstractmethod + def list(self, prefix: str) -> List[str]: + raise NotImplementedError() + + +class BaseFileConnector(BaseConnector): + """ + List full file names from remote fs path and filter by allow pattern. + + Args: + allow_pattern: A list of patterns of which files to pull. + + Returns: + list[str]: List of full paths allowed by the pattern + """ + + @abstractmethod + def glob(self, allow_pattern: str) -> List[str]: + raise NotImplementedError() diff --git a/python/sglang/srt/connector/redis.py b/python/sglang/srt/connector/redis.py new file mode 100644 index 000000000..cb1db3f7c --- /dev/null +++ b/python/sglang/srt/connector/redis.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 + +import logging +from typing import Generator, List, Optional, Tuple +from urllib.parse import urlparse + +import torch + +from sglang.srt.connector import BaseKVConnector +from sglang.srt.connector.serde import create_serde +from sglang.srt.connector.utils import pull_files_from_db + +logger = logging.getLogger(__name__) + + +class RedisConnector(BaseKVConnector): + + def __init__(self, url: str): + import redis + + super().__init__(url) + parsed_url = urlparse(url) + self.connection = redis.Redis(host=parsed_url.hostname, port=parsed_url.port) + self.model_name = parsed_url.path.lstrip("/") + # TODO: more serde options + self.s, self.d = create_serde("safe") + + def get(self, key: str) -> Optional[torch.Tensor]: + val = self.connection.get(key) + + if val is None: + logger.error("Key %s not found", key) + return None + + return self.d.from_bytes(val) + + def getstr(self, key: str) -> Optional[str]: + val = self.connection.get(key) + if val is None: + logger.error("Key %s not found", key) + return None + + return val.decode("utf-8") + + def set(self, key: str, tensor: torch.Tensor) -> None: + assert tensor is not None + self.connection.set(key, self.s.to_bytes(tensor)) + + def setstr(self, key: str, obj: str) -> None: + self.connection.set(key, obj) + + def list(self, prefix: str) -> List[str]: + cursor = 0 + all_keys: List[bytes] = [] + + while True: + ret: Tuple[int, List[bytes]] = self.connection.scan( + cursor=cursor, match=f"{prefix}*" + ) # type: ignore + cursor, keys = ret + all_keys.extend(keys) + if cursor == 0: + break + + return [key.decode("utf-8") for key in all_keys] + + def weight_iterator( + self, rank: int = 0 + ) -> Generator[Tuple[str, bytes], None, None]: + keys = self.list(f"{self.model_name}/keys/rank_{rank}/") + for key in keys: + val = self.get(key) + key = key.removeprefix(f"{self.model_name}/keys/rank_{rank}/") + yield key, val + + def pull_files( + self, + allow_pattern: Optional[List[str]] = None, + ignore_pattern: Optional[List[str]] = None, + ) -> None: + pull_files_from_db(self, self.model_name, allow_pattern, ignore_pattern) + + def close(self): + self.connection.close() + super().close() diff --git a/python/sglang/srt/connector/s3.py b/python/sglang/srt/connector/s3.py new file mode 100644 index 000000000..7bef8f5d5 --- /dev/null +++ b/python/sglang/srt/connector/s3.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 + +import fnmatch +import os +from pathlib import Path +from typing import Generator, Optional, Tuple + +import torch + +from sglang.srt.connector import BaseFileConnector + + +def _filter_allow(paths: list[str], patterns: list[str]) -> list[str]: + return [ + path + for path in paths + if any(fnmatch.fnmatch(path, pattern) for pattern in patterns) + ] + + +def _filter_ignore(paths: list[str], patterns: list[str]) -> list[str]: + return [ + path + for path in paths + if not any(fnmatch.fnmatch(path, pattern) for pattern in patterns) + ] + + +def list_files( + s3, + path: str, + allow_pattern: Optional[list[str]] = None, + ignore_pattern: Optional[list[str]] = None, +) -> tuple[str, str, list[str]]: + """ + List files from S3 path and filter by pattern. + + Args: + s3: S3 client to use. + path: The S3 path to list from. + allow_pattern: A list of patterns of which files to pull. + ignore_pattern: A list of patterns of which files not to pull. + + Returns: + tuple[str, str, list[str]]: A tuple where: + - The first element is the bucket name + - The second element is string represent the bucket + and the prefix as a dir like string + - The third element is a list of files allowed or + disallowed by pattern + """ + parts = path.removeprefix("s3://").split("/") + prefix = "/".join(parts[1:]) + bucket_name = parts[0] + + objects = s3.list_objects_v2(Bucket=bucket_name, Prefix=prefix) + paths = [obj["Key"] for obj in objects.get("Contents", [])] + + paths = _filter_ignore(paths, ["*/"]) + if allow_pattern is not None: + paths = _filter_allow(paths, allow_pattern) + + if ignore_pattern is not None: + paths = _filter_ignore(paths, ignore_pattern) + + return bucket_name, prefix, paths + + +class S3Connector(BaseFileConnector): + + def __init__(self, url: str) -> None: + import boto3 + + super().__init__(url) + self.client = boto3.client("s3") + + def glob(self, allow_pattern: Optional[list[str]] = None) -> list[str]: + bucket_name, _, paths = list_files( + self.client, path=self.url, allow_pattern=allow_pattern + ) + return [f"s3://{bucket_name}/{path}" for path in paths] + + def pull_files( + self, + allow_pattern: Optional[list[str]] = None, + ignore_pattern: Optional[list[str]] = None, + ) -> None: + """ + Pull files from S3 storage into the temporary directory. + + Args: + s3_model_path: The S3 path of the model. + allow_pattern: A list of patterns of which files to pull. + ignore_pattern: A list of patterns of which files not to pull. + + """ + bucket_name, base_dir, files = list_files( + self.client, self.url, allow_pattern, ignore_pattern + ) + if len(files) == 0: + return + + for file in files: + destination_file = os.path.join(self.local_dir, file.removeprefix(base_dir)) + local_dir = Path(destination_file).parent + os.makedirs(local_dir, exist_ok=True) + self.client.download_file(bucket_name, file, destination_file) + + def weight_iterator( + self, rank: int = 0 + ) -> Generator[Tuple[str, torch.Tensor], None, None]: + from sglang.srt.model_loader.weight_utils import ( + runai_safetensors_weights_iterator, + ) + + # only support safetensor files now + hf_weights_files = self.glob(allow_pattern=["*.safetensors"]) + return runai_safetensors_weights_iterator(hf_weights_files) + + def close(self): + self.client.close() + super().close() diff --git a/python/sglang/srt/connector/serde/__init__.py b/python/sglang/srt/connector/serde/__init__.py new file mode 100644 index 000000000..c05b20afa --- /dev/null +++ b/python/sglang/srt/connector/serde/__init__.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 + +# inspired by LMCache +from typing import Optional, Tuple + +import torch + +from sglang.srt.connector.serde.safe_serde import SafeDeserializer, SafeSerializer +from sglang.srt.connector.serde.serde import Deserializer, Serializer + + +def create_serde(serde_type: str) -> Tuple[Serializer, Deserializer]: + s: Optional[Serializer] = None + d: Optional[Deserializer] = None + + if serde_type == "safe": + s = SafeSerializer() + d = SafeDeserializer() + else: + raise ValueError(f"Unknown serde type: {serde_type}") + + return s, d + + +__all__ = [ + "Serializer", + "Deserializer", + "SafeSerializer", + "SafeDeserializer", + "create_serde", +] diff --git a/python/sglang/srt/connector/serde/safe_serde.py b/python/sglang/srt/connector/serde/safe_serde.py new file mode 100644 index 000000000..3e75f9bfc --- /dev/null +++ b/python/sglang/srt/connector/serde/safe_serde.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import Union + +import torch +from safetensors.torch import load, save + +from sglang.srt.connector.serde.serde import Deserializer, Serializer + + +class SafeSerializer(Serializer): + + def __init__(self): + super().__init__() + + def to_bytes(self, t: torch.Tensor) -> bytes: + return save({"tensor_bytes": t.cpu().contiguous()}) + + +class SafeDeserializer(Deserializer): + + def __init__(self): + # TODO: dtype options + super().__init__(torch.float32) + + def from_bytes_normal(self, b: Union[bytearray, bytes]) -> torch.Tensor: + return load(bytes(b))["tensor_bytes"] + + def from_bytes(self, b: Union[bytearray, bytes]) -> torch.Tensor: + return self.from_bytes_normal(b) diff --git a/python/sglang/srt/connector/serde/serde.py b/python/sglang/srt/connector/serde/serde.py new file mode 100644 index 000000000..3d6f804d7 --- /dev/null +++ b/python/sglang/srt/connector/serde/serde.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 + +import abc +from abc import ABC, abstractmethod + +import torch + + +class Serializer(ABC): + + @abstractmethod + def to_bytes(self, t: torch.Tensor) -> bytes: + """ + Serialize a pytorch tensor to bytes. The serialized bytes should contain + both the data and the metadata (shape, dtype, etc.) of the tensor. + + Input: + t: the input pytorch tensor, can be on any device, in any shape, + with any dtype + + Returns: + bytes: the serialized bytes + """ + raise NotImplementedError + + +class Deserializer(metaclass=abc.ABCMeta): + + def __init__(self, dtype): + self.dtype = dtype + + @abstractmethod + def from_bytes(self, bs: bytes) -> torch.Tensor: + """ + Deserialize a pytorch tensor from bytes. + + Input: + bytes: a stream of bytes + + Output: + torch.Tensor: the deserialized pytorch tensor + """ + raise NotImplementedError diff --git a/python/sglang/srt/connector/utils.py b/python/sglang/srt/connector/utils.py new file mode 100644 index 000000000..6cd12da33 --- /dev/null +++ b/python/sglang/srt/connector/utils.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +from pathlib import Path +from typing import Optional +from urllib.parse import urlparse + +from sglang.srt.connector import BaseConnector + + +def parse_model_name(url: str) -> str: + """ + Parse the model name from the url. + Only used for db connector + """ + parsed_url = urlparse(url) + return parsed_url.path.lstrip("/") + + +def pull_files_from_db( + connector: BaseConnector, + model_name: str, + allow_pattern: Optional[list[str]] = None, + ignore_pattern: Optional[list[str]] = None, +) -> None: + prefix = f"{model_name}/files/" + local_dir = connector.get_local_dir() + files = connector.list(prefix) + + for file in files: + destination_file = os.path.join(local_dir, file.removeprefix(prefix)) + local_dir = Path(destination_file).parent + os.makedirs(local_dir, exist_ok=True) + with open(destination_file, "wb") as f: + f.write(connector.getstr(file).encode("utf-8")) diff --git a/python/sglang/srt/constants.py b/python/sglang/srt/constants.py new file mode 100644 index 000000000..aa03a089b --- /dev/null +++ b/python/sglang/srt/constants.py @@ -0,0 +1,3 @@ +# GPU Memory Types +GPU_MEMORY_TYPE_KV_CACHE = "kv_cache" +GPU_MEMORY_TYPE_WEIGHTS = "weights" diff --git a/python/sglang/srt/constrained/base_grammar_backend.py b/python/sglang/srt/constrained/base_grammar_backend.py new file mode 100644 index 000000000..4fe5d6c77 --- /dev/null +++ b/python/sglang/srt/constrained/base_grammar_backend.py @@ -0,0 +1,213 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""The baseclass of a backend for grammar-guided constrained decoding.""" + +import logging +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from threading import Event +from typing import Dict, List, Optional, Tuple + +import torch + +from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +class BaseGrammarObject: + + def __init__(self): + self._finished = False + + def accept_token(self, token: int) -> None: + """ + Accept a token in the grammar. + """ + raise NotImplementedError() + + def rollback(self, k: int): + raise NotImplementedError() + + def is_terminated(self): + return False + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device + ) -> torch.Tensor: + raise NotImplementedError() + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + raise NotImplementedError() + + @staticmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor: + raise NotImplementedError() + + @staticmethod + def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + raise NotImplementedError() + + def copy(self) -> "BaseGrammarObject": + return self + + @property + def finished(self): + return self._finished + + @finished.setter + def finished(self, finished): + self._finished = finished + + def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]: + """ + Try to jump forward in the grammar. + + Returns: + A jump forward helper which may be used in `jump_forward_str_state`. + None if the jump forward is not possible. + """ + raise NotImplementedError() + + def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]: + """ + Jump forward for the grammar. + + Returns: + A tuple of the jump forward string and the next state of the grammar + (which can be used in `jump_and_retokenize` if needed). + """ + raise NotImplementedError() + + def jump_and_retokenize( + self, old_output_ids: List[int], new_output_ids: List[int], next_state: int + ) -> None: + """ + Jump forward occurs, and update the grammar state if needed. + """ + raise NotImplementedError() + + +INVALID_GRAMMAR_OBJ = BaseGrammarObject() + + +@dataclass +class CacheEntry: + value: BaseGrammarObject + event: Event + + +class BaseGrammarBackend: + def __init__(self): + self.executor = ThreadPoolExecutor() + self.cache: Dict[Tuple[str, str], CacheEntry] = {} + + def _not_supported(self, key_type: str, key_string: str) -> None: + logger.warning(f"Skip unsupported {key_type=}, {key_string=}") + + def dispatch_fallback( + self, key_type: str, key_string: str + ) -> Optional[BaseGrammarObject]: + """ + This function should not be reached in any case. + """ + raise ValueError(f"Invalid key_type: {key_type}={key_string}") + + def dispatch_json(self, key_string: str) -> Optional[BaseGrammarObject]: + return self._not_supported("json", key_string) + + def dispatch_regex(self, key_string: str) -> Optional[BaseGrammarObject]: + return self._not_supported("regex", key_string) + + def dispatch_ebnf(self, key_string: str) -> Optional[BaseGrammarObject]: + return self._not_supported("ebnf", key_string) + + def dispatch_structural_tag(self, key_string: str) -> Optional[BaseGrammarObject]: + return self._not_supported("structural_tag", key_string) + + def _init_value_dispatch(self, key: Tuple[str, str]) -> Optional[BaseGrammarObject]: + key_type, key_string = key + if key_type == "json": + return self.dispatch_json(key_string) + elif key_type == "regex": + return self.dispatch_regex(key_string) + elif key_type == "ebnf": + return self.dispatch_ebnf(key_string) + elif key_type == "structural_tag": + return self.dispatch_structural_tag(key_string) + elif key_type == "structural_pattern": + return self.dispatch_structural_pattern(key_string) + else: + return self.dispatch_fallback(key_type, key_string) + + def get_cached_or_future_value( + self, key: Tuple[str, str] + ) -> Optional[BaseGrammarObject]: + value = self.cache.get(key) + if value: + return value.copy(), True + value = self.executor.submit(self._init_value_dispatch, key) + return value, False + + def set_cache(self, key: Tuple[str, str], value: BaseGrammarObject): + self.cache[key] = value + + def reset(self): + self.cache.clear() + + +def create_grammar_backend( + server_args: ServerArgs, + tokenizer, + vocab_size: int, + eos_token_ids: Optional[set] = None, +) -> Optional[BaseGrammarBackend]: + if server_args.grammar_backend == "outlines": + from sglang.srt.constrained.outlines_backend import OutlinesGrammarBackend + + grammar_backend = OutlinesGrammarBackend( + tokenizer, + whitespace_pattern=server_args.constrained_json_whitespace_pattern, + ) + elif server_args.grammar_backend == "xgrammar": + from sglang.srt.constrained.xgrammar_backend import XGrammarGrammarBackend + + # Convert Set[int] to List[int] if needed + eos_list = list(eos_token_ids) if eos_token_ids else None + + grammar_backend = XGrammarGrammarBackend( + tokenizer, vocab_size=vocab_size, model_eos_token_ids=eos_list + ) + elif server_args.grammar_backend == "llguidance": + from sglang.srt.constrained.llguidance_backend import GuidanceBackend + + grammar_backend = GuidanceBackend( + tokenizer=tokenizer, + whitespace_pattern=server_args.constrained_json_whitespace_pattern, + ) + elif server_args.grammar_backend == "none": + return None + else: + raise ValueError(f"Invalid grammar backend: {server_args.grammar_backend}") + + if server_args.reasoning_parser and hasattr(tokenizer, "think_end_id"): + from sglang.srt.constrained.reasoner_grammar_backend import ( + ReasonerGrammarBackend, + ) + + grammar_backend = ReasonerGrammarBackend( + grammar_backend, tokenizer.think_end_id + ) + + return grammar_backend diff --git a/python/sglang/srt/constrained/llguidance_backend.py b/python/sglang/srt/constrained/llguidance_backend.py new file mode 100644 index 000000000..2acbf2c51 --- /dev/null +++ b/python/sglang/srt/constrained/llguidance_backend.py @@ -0,0 +1,174 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Constrained decoding with llguidance backend.""" + +import json +import logging +import os +from typing import List, Optional, Tuple + +import torch +from llguidance import LLMatcher, LLTokenizer, StructTag, grammar_from +from llguidance.hf import from_tokenizer +from llguidance.torch import ( + allocate_token_bitmask, + apply_token_bitmask_inplace, + fill_next_token_bitmask, +) + +from sglang.srt.constrained.base_grammar_backend import ( + INVALID_GRAMMAR_OBJ, + BaseGrammarBackend, + BaseGrammarObject, +) + +logger = logging.getLogger(__name__) + + +class GuidanceGrammar(BaseGrammarObject): + + def __init__(self, llguidance_tokenizer: LLTokenizer, serialized_grammar: str): + super().__init__() + self.llguidance_tokenizer = llguidance_tokenizer + self.serialized_grammar = serialized_grammar + + self.ll_matcher = LLMatcher( + self.llguidance_tokenizer, + self.serialized_grammar, + log_level=int(os.environ.get("LLGUIDANCE_LOG_LEVEL", "1")), + ) + self.finished = False + self.bitmask = None + + def accept_token(self, token: int): + if not self.ll_matcher.consume_token(token): + logger.warning(f"matcher error: {self.ll_matcher.get_error()}") + self.finished = True + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + if self.ll_matcher.is_stopped(): + self.finished = True + + fill_next_token_bitmask(self.ll_matcher, vocab_mask, idx) + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device + ) -> torch.Tensor: + if self.bitmask is None or self.bitmask.shape[0] < batch_size: + # only create bitmask when batch gets larger + self.bitmask = allocate_token_bitmask( + batch_size, self.llguidance_tokenizer.vocab_size + ) + bitmask = self.bitmask + else: + bitmask = self.bitmask[:batch_size] + + return bitmask + + @staticmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor: + return vocab_mask.to(device, non_blocking=True) + + @staticmethod + def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + apply_token_bitmask_inplace(logits, vocab_mask) + + def copy(self): + return GuidanceGrammar( + llguidance_tokenizer=self.llguidance_tokenizer, + serialized_grammar=self.serialized_grammar, + ) + + def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]: + ff_tokens = self.ll_matcher.compute_ff_tokens() + if ff_tokens: + return ff_tokens, "" + else: + return None + + def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]: + return "", -1 + + def jump_and_retokenize( + self, old_output_ids: List[int], new_output_ids: List[int], next_state: int + ): + pass + + +class GuidanceBackend(BaseGrammarBackend): + + def __init__( + self, + tokenizer, + whitespace_pattern: Optional[str] = None, + n_vocab: Optional[int] = None, + ): + super().__init__() + + self.tokenizer = tokenizer + self.whitespace_pattern = whitespace_pattern + self.llguidance_tokenizer = from_tokenizer(self.tokenizer, n_vocab) + + def _from_serialized(self, serialized_grammar) -> Optional[GuidanceGrammar]: + try: + return GuidanceGrammar( + llguidance_tokenizer=self.llguidance_tokenizer, + serialized_grammar=serialized_grammar, + ) + except Exception as e: + logger.error(f"Hit invalid grammar: {serialized_grammar=}, {e=}") + return INVALID_GRAMMAR_OBJ + + def dispatch_json(self, key_string: str) -> Optional[GuidanceGrammar]: + try: + serialized_grammar = LLMatcher.grammar_from_json_schema( + key_string, + defaults={ + "whitespace_pattern": self.whitespace_pattern, + }, + ) + except Exception as e: + logger.error(f"Hit invalid json_schema: {key_string=}, {e=}") + return INVALID_GRAMMAR_OBJ + return self._from_serialized(serialized_grammar) + + def dispatch_regex(self, key_string: str) -> Optional[GuidanceGrammar]: + serialized_grammar = grammar_from("regex", key_string) + return self._from_serialized(serialized_grammar) + + def dispatch_ebnf(self, key_string: str) -> Optional[GuidanceGrammar]: + try: + serialized_grammar = grammar_from("ebnf", key_string) + return self._from_serialized(serialized_grammar) + except ValueError as e: + logger.error(f"Hit invalid ebnf: {key_string=}, {e=}") + return INVALID_GRAMMAR_OBJ + + def dispatch_structural_tag(self, key_string: str) -> Optional[GuidanceGrammar]: + try: + structural_tag = json.loads(key_string) + tags = [ + StructTag( + begin=structure["begin"], + grammar=structure["schema"], + end=structure["end"], + trigger=structural_tag["triggers"][0], # TODO? + ) + for structure in structural_tag["structures"] + ] + g = StructTag.to_grammar(tags) + return self._from_serialized(g) + except Exception as e: + logging.error(f"Hit invalid structural_tag: {key_string=}, {e=}") + return INVALID_GRAMMAR_OBJ diff --git a/python/sglang/srt/constrained/outlines_backend.py b/python/sglang/srt/constrained/outlines_backend.py new file mode 100644 index 000000000..5302fadaa --- /dev/null +++ b/python/sglang/srt/constrained/outlines_backend.py @@ -0,0 +1,191 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Constrained decoding with outlines backend.""" + +import json +import logging +from typing import Dict, List, Optional, Tuple, Union + +import interegular +import torch +from outlines.fsm.guide import RegexGuide +from outlines.models.transformers import TransformerTokenizer +from pydantic import BaseModel + +from sglang.srt.constrained.base_grammar_backend import ( + INVALID_GRAMMAR_OBJ, + BaseGrammarBackend, + BaseGrammarObject, +) +from sglang.srt.constrained.outlines_jump_forward import OutlinesJumpForwardMap + +try: + from outlines.fsm.json_schema import build_regex_from_schema +except ImportError: + from outlines_core.fsm.json_schema import build_regex_from_schema + + +logger = logging.getLogger(__name__) + + +class OutlinesGrammar(BaseGrammarObject): + def __init__( + self, + guide: RegexGuide, + jump_forward_map: Union[OutlinesJumpForwardMap, None], + ) -> None: + super().__init__() + self.guide = guide + self.jump_forward_map = jump_forward_map + self.state = 0 + self.finished = False + + def accept_token(self, token: int): + self.state = self.guide.get_next_state(self.state, token) + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device + ) -> torch.Tensor: + return torch.zeros(batch_size, vocab_size, dtype=torch.bool, device=device) + + @staticmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor: + return vocab_mask + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + tokens = torch.tensor( + self.guide.get_next_instruction(self.state).tokens, dtype=torch.int64 + ).to(vocab_mask.device, non_blocking=True) + vocab_mask = vocab_mask[idx] + vocab_mask.fill_(1) + vocab_mask.scatter_(0, tokens, torch.zeros_like(tokens, dtype=torch.bool)) + + @staticmethod + def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor): + logits.masked_fill_(vocab_mask, float("-inf")) + + def copy(self): + return OutlinesGrammar(self.guide, self.jump_forward_map) + + def try_jump_forward(self, tokenizer) -> Optional[Tuple]: + if not self.jump_forward_map: + return None + + jump_forward_bytes = self.jump_forward_map.jump_forward_byte(self.state) + if jump_forward_bytes is None or len(jump_forward_bytes) <= 1: + return None + + # preprocess the jump forward string + suffix_bytes = [] + continuation_range = range(0x80, 0xC0) + cur_state = self.state + while ( + len(jump_forward_bytes) and jump_forward_bytes[0][0] in continuation_range + ): + # continuation bytes + byte_edge = jump_forward_bytes.pop(0) + suffix_bytes.append(byte_edge[0]) + cur_state = byte_edge[1] + + suffix_tokens = [f"<0x{hex(b)[2:].upper()}>" for b in suffix_bytes] + suffix_ids = tokenizer.convert_tokens_to_ids(suffix_tokens) + return suffix_ids, cur_state + + def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]: + _, cur_state = helper + return self.jump_forward_map.jump_forward_symbol(cur_state) + + def jump_and_retokenize( + self, old_output_ids: List[int], new_output_ids: List[int], next_state: int + ): + self.state = next_state + + +class OutlinesGrammarBackend(BaseGrammarBackend): + def __init__( + self, + tokenizer, + whitespace_pattern: bool, + ): + super().__init__() + + try: + self.outlines_tokenizer = TransformerTokenizer(tokenizer) + except AttributeError: + # FIXME: tmp fix for chatglm2 & chatglm3 (pad_token_id=0) + origin_pad_token_id = tokenizer.pad_token_id + + def fset(self, value): + self._value = value + + type(tokenizer).pad_token_id = property( + fget=type(tokenizer).pad_token_id.fget, fset=fset + ) + self.outlines_tokenizer = TransformerTokenizer(tokenizer) + self.outlines_tokenizer.tokenizer.pad_token_id = origin_pad_token_id + self.outlines_tokenizer.pad_token_id = origin_pad_token_id + self.outlines_tokenizer.pad_token = ( + self.outlines_tokenizer.tokenizer.pad_token + ) + self.outlines_tokenizer.vocabulary = ( + self.outlines_tokenizer.tokenizer.get_vocab() + ) + self.whitespace_pattern = whitespace_pattern + + def _compile_regex(self, regex: str) -> Optional[OutlinesGrammar]: + try: + if hasattr(RegexGuide, "from_regex"): + # outlines >= 0.1.1 + guide = RegexGuide.from_regex(regex, self.outlines_tokenizer) + else: + # outlines <= 0.0.46 + guide = RegexGuide(regex, self.outlines_tokenizer) + except interegular.patterns.InvalidSyntax as e: + logger.error(f"Hit invalid regex schema: {regex=}, {e=}") + return INVALID_GRAMMAR_OBJ + + jump_forward_map = None + return OutlinesGrammar(guide, jump_forward_map) + + def dispatch_ebnf(self, key_string: str): + return super().dispatch_ebnf(key_string) + + def dispatch_structural_tag(self, key_string: str): + return super().dispatch_structural_tag(key_string) + + def dispatch_json(self, key_string: str): + try: + regex = build_regex_from_object( + key_string, + whitespace_pattern=self.whitespace_pattern, + ) + except (NotImplementedError, json.decoder.JSONDecodeError, ValueError) as e: + logger.error(f"Hit invalid json_schema: {key_string=}, {e=}") + return INVALID_GRAMMAR_OBJ + return self._compile_regex(regex) + + def dispatch_regex(self, key_string: str): + return self._compile_regex(key_string) + + +def build_regex_from_object( + object: Union[str, BaseModel, Dict], whitespace_pattern: Optional[str] = None +): + if isinstance(object, type(BaseModel)): + schema = json.dumps(object.model_json_schema()) + elif isinstance(object, Dict): + schema = json.dumps(object) + else: + schema = object + return build_regex_from_schema(schema, whitespace_pattern) diff --git a/python/sglang/srt/constrained/outlines_jump_forward.py b/python/sglang/srt/constrained/outlines_jump_forward.py new file mode 100644 index 000000000..cfc65f75f --- /dev/null +++ b/python/sglang/srt/constrained/outlines_jump_forward.py @@ -0,0 +1,200 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +""" +Faster constrained decoding with jump forward decoding / compressed finite state machine. +Reference: https://lmsys.org/blog/2024-02-05-compressed-fsm/ +""" + +import dataclasses +import logging +from collections import defaultdict +from typing import Optional + +import interegular +from interegular import InvalidSyntax +from outlines.caching import cache + +from sglang.srt.utils import get_bool_env_var + +try: + # outlines >= 0.1.0 + from outlines_core.fsm.outlines_core_rs import FSMInfo + from outlines_core.fsm.regex import make_byte_level_fsm, make_deterministic_fsm +except ImportError: + # outlines <= 0.0.46 + from outlines.fsm.regex import FSMInfo, make_byte_level_fsm, make_deterministic_fsm + +IP_REGEX = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)" + +# Env var was set in sglang.srt.server_args.ServerArgs.__post__init__ +DISABLE_DISK_CACHE = get_bool_env_var("SGLANG_DISABLE_OUTLINES_DISK_CACHE", "true") + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class JumpEdge: + symbol: str = None + symbol_next_state: int = None + byte: int = None + byte_next_state: int = None + + +def disk_cache(expire: Optional[float] = None, typed=False, ignore=()): + if not DISABLE_DISK_CACHE: + return cache(expire, typed, ignore) + else: + return lambda fn: None + + +@disk_cache() +def init_state_to_jump_forward(regex_string): + try: + regex_pattern = interegular.parse_pattern(regex_string) + except InvalidSyntax as e: + logger.warning(f"skip invalid regex: {regex_string}, {e=}") + return + + byte_fsm = make_byte_level_fsm(regex_pattern.to_fsm().reduce(), keep_utf8=True) + regex_fsm, _ = make_deterministic_fsm(byte_fsm) + + fsm_info: FSMInfo = regex_fsm.fsm_info + + symbol_to_id = fsm_info.alphabet_symbol_mapping + id_to_symbol = {} + for symbol, id_ in symbol_to_id.items(): + id_to_symbol.setdefault(id_, []).append(symbol) + + transitions = fsm_info.transitions + + outgoings_ct = defaultdict(int) + # NOTE(lsyin): Final states can lead to terminate, so they have one outgoing edge naturally + for s in fsm_info.finals: + outgoings_ct[s] = 1 + + state_to_jump_forward = {} + for (state, id_), next_state in transitions.items(): + if id_ == fsm_info.alphabet_anything_value: + # Arbitrarily symbol cannot be recognized as jump forward + continue + + symbols = id_to_symbol[id_] + for c in symbols: + if len(c) > 1: + # Skip byte level transitions like c = "5E" + continue + + outgoings_ct[state] += 1 + if outgoings_ct[state] > 1: + if state in state_to_jump_forward: + del state_to_jump_forward[state] + break + + state_to_jump_forward[state] = JumpEdge( + symbol=c, + symbol_next_state=next_state, + ) + + # Process the byte level jump forward + outgoings_ct = defaultdict(int) + for s in fsm_info.finals: + outgoings_ct[s] = 1 + + for (state, id_), next_state in transitions.items(): + if id_ == fsm_info.alphabet_anything_value: + continue + symbols = id_to_symbol[id_] + for c in symbols: + byte_ = None + if len(c) == 1 and ord(c) < 0x80: + # ASCII character + byte_ = ord(c) + elif len(c) > 1: + # FIXME: This logic is due to the leading \x00 + # https://github.com/outlines-dev/outlines/pull/930 + byte_ = int(symbols[0][1:], 16) + + if byte_ is not None: + outgoings_ct[state] += 1 + if outgoings_ct[state] > 1: + if state in state_to_jump_forward: + del state_to_jump_forward[state] + break + e = state_to_jump_forward.get(state, JumpEdge()) + e.byte = byte_ + e.byte_next_state = next_state + state_to_jump_forward[state] = e + + return state_to_jump_forward + + +class OutlinesJumpForwardMap: + def __init__(self, regex_string): + self.state_to_jump_forward = init_state_to_jump_forward(regex_string) + + def jump_forward_symbol(self, state): + jump_forward_str = "" + next_state = state + while state in self.state_to_jump_forward: + e = self.state_to_jump_forward[state] + if e.symbol is None: + break + jump_forward_str += e.symbol + next_state = e.symbol_next_state + state = next_state + + return jump_forward_str, next_state + + def jump_forward_byte(self, state): + if state not in self.state_to_jump_forward: + return None + + jump_forward_bytes = [] + next_state = None + while state in self.state_to_jump_forward: + e = self.state_to_jump_forward[state] + assert e.byte is not None and e.byte_next_state is not None + jump_forward_bytes.append((e.byte, e.byte_next_state)) + next_state = e.byte_next_state + state = next_state + + return jump_forward_bytes + + def is_jump_forward_symbol_state(self, state): + return ( + state in self.state_to_jump_forward + and self.state_to_jump_forward[state].symbol is not None + ) + + +def test_main(regex_string): + jump_forward_map = OutlinesJumpForwardMap(regex_string) + for state, e in jump_forward_map.state_to_jump_forward.items(): + if e.symbol is not None: + jump_forward_str, next_state = jump_forward_map.jump_forward_symbol(state) + print(f"{state} -> {next_state}", jump_forward_str) + bytes_ = jump_forward_map.jump_forward_byte(state) + print(f"{state} -> {bytes_[-1][1]}", [hex(b) for b, _ in bytes_]) + + +if __name__ == "__main__": + import outlines + + outlines.caching.clear_cache() + test_main(r"The google's DNS sever address is " + IP_REGEX) + test_main(r"霍格沃茨特快列车|霍比特人比尔博") + # 霍格: \xe9\x9c\x8d \xe6\xa0\xbc ... + # 霍比: \xe9\x9c\x8d \xe6\xaf\x94 ... + + test_main(r"[-+]?[0-9]+[ ]*") diff --git a/python/sglang/srt/constrained/reasoner_grammar_backend.py b/python/sglang/srt/constrained/reasoner_grammar_backend.py new file mode 100644 index 000000000..ca5f118ef --- /dev/null +++ b/python/sglang/srt/constrained/reasoner_grammar_backend.py @@ -0,0 +1,90 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""The baseclass of a backend for reasoner grammar-guided constrained decoding.""" + +from typing import List, Optional, Tuple + +import torch + +from .base_grammar_backend import BaseGrammarBackend, BaseGrammarObject + + +class ReasonerGrammarObject(BaseGrammarObject): + def __init__(self, grammar: BaseGrammarObject, think_end_id): + super().__init__() + self.grammar = grammar + self.think_end_id = think_end_id + self.is_in_reasoning = True + + def accept_token(self, token: int): + if token == self.think_end_id: + self.is_in_reasoning = False + + if not self.is_in_reasoning and token != self.think_end_id: + self.grammar.accept_token(token) + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device + ) -> torch.Tensor: + return self.grammar.allocate_vocab_mask(vocab_size, batch_size, device) + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + if not self.is_in_reasoning: + self.grammar.fill_vocab_mask(vocab_mask, idx) + + def move_vocab_mask(self, vocab_mask: torch.Tensor, device) -> torch.Tensor: + return self.grammar.move_vocab_mask(vocab_mask, device) + + @property + def apply_vocab_mask(self): + return self.grammar.apply_vocab_mask + + def copy(self) -> BaseGrammarObject: + return ReasonerGrammarObject(self.grammar.copy(), self.think_end_id) + + @property + def finished(self): + return self.grammar.finished + + @finished.setter + def finished(self, finished): + self.grammar.finished = finished + + def try_jump_forward(self, tokenizer): + return self.grammar.try_jump_forward(tokenizer) + + def jump_forward_str_state(self, helper): + return self.grammar.jump_forward_str_state(helper) + + def jump_and_retokenize( + self, old_output_ids: List[int], new_output_ids: List[int], next_state: int + ): + return self.grammar.jump_and_retokenize( + old_output_ids, new_output_ids, next_state + ) + + +class ReasonerGrammarBackend(BaseGrammarBackend): + def __init__(self, grammar_backend: BaseGrammarBackend, think_end_id): + super().__init__() + self.grammar_backend = grammar_backend + self.think_end_id = think_end_id + + def _init_value_dispatch( + self, key: Tuple[str, str] + ) -> Optional[ReasonerGrammarObject]: + ret = self.grammar_backend._init_value_dispatch(key) + if ret is None: + return None + return ReasonerGrammarObject(ret, self.think_end_id) diff --git a/python/sglang/srt/constrained/triton_ops/bitmask_ops.py b/python/sglang/srt/constrained/triton_ops/bitmask_ops.py new file mode 100644 index 000000000..9a195c006 --- /dev/null +++ b/python/sglang/srt/constrained/triton_ops/bitmask_ops.py @@ -0,0 +1,141 @@ +# Adapt from +# https://github.com/mlc-ai/xgrammar/blob/v0.1.17/python/xgrammar/kernels/apply_token_bitmask_inplace_triton.py + +from typing import List, Optional, Union + +import torch +import triton +import triton.language as tl + +from sglang.srt.utils import get_device_core_count + + +@triton.jit +def apply_token_bitmask_inplace_kernel( + logits_ptr, + bitmask_ptr, + indices_ptr, + num_rows, + vocab_size, + logits_strides, + bitmask_strides, + NUM_SMS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Apply a bitmask to logits in-place using Triton. The bitmask is a 01 bitwise compressed tensor, + where 0 means the token is masked and 1 means the token is not masked. After applying the bitmask, + the masked logits will be set to -inf. + + Parameters + ---------- + logits_ptr : tl.tensor + Pointer to the logits tensor to apply the bitmask to. + + bitmask_ptr : tl.tensor + Pointer to the bitmask tensor to apply. + + indices_ptr : Optional[tl.tensor] + Optional pointer to indices tensor specifying which rows to apply the mask to. + + num_rows : int + Number of rows to process. If indices_ptr is provided, this is the number of unique indices. + + vocab_size : int + Size of the vocabulary dimension. If the logits does not have a vocab padding, this is the + same as the logits's second dimension. Otherwise, this is the actual size of the vocabulary. + + logits_strides : int + Stride between rows in the logits tensor. + + bitmask_strides : int + Stride between rows in the bitmask tensor. + + NUM_SMS : int + Number of streaming multiprocessors to use. + + BLOCK_SIZE : int + Size of processing blocks. + """ + + pid = tl.program_id(0) + num_blocks = tl.cdiv(vocab_size, BLOCK_SIZE) + for work_id in tl.range(pid, num_rows * num_blocks, NUM_SMS): + row_id = work_id // num_blocks + block_offset = (work_id % num_blocks) * BLOCK_SIZE + batch_id = row_id if indices_ptr is None else tl.load(indices_ptr + row_id) + offsets = block_offset + tl.arange(0, BLOCK_SIZE) + bitmask_offsets = block_offset // 32 + tl.arange(0, BLOCK_SIZE // 32) + vocab_mask = offsets < vocab_size + packed_bitmask_mask = bitmask_offsets < bitmask_strides + packed_bitmask = tl.load( + bitmask_ptr + batch_id * bitmask_strides + bitmask_offsets, + packed_bitmask_mask, + ) + bitmask = ((packed_bitmask[:, None] >> (tl.arange(0, 32)[None, :])) & 1) == 0 + bitmask = bitmask.reshape(BLOCK_SIZE) + + tl.store( + logits_ptr + batch_id * logits_strides + offsets, + -float("inf"), + vocab_mask & bitmask, + ) + + +def apply_token_bitmask_inplace_triton( + logits: torch.Tensor, + bitmask: torch.Tensor, + indices: Optional[Union[List[int], torch.Tensor]] = None, +): + NUM_SMS = get_device_core_count() + BLOCK_SIZE = 4096 + BITS_PER_BLOCK = 32 + + # Check input dtype + assert bitmask.dtype == torch.int32, "bitmask must be of type int32" + + # Check input tensor shapes. + logits_shape = logits.shape + bitmask_shape = bitmask.shape + if logits.ndim == 1: + logits_shape = (1, logits_shape[0]) + if bitmask.ndim == 1: + bitmask_shape = (1, bitmask_shape[0]) + + required_bitmask_width = (logits_shape[1] + BITS_PER_BLOCK - 1) // BITS_PER_BLOCK + assert required_bitmask_width >= bitmask_shape[1], ( + f"Bitmask width too large: allow at most {required_bitmask_width} int32s for " + f"logits' width {logits_shape[1]}, but got {bitmask_shape[1]}" + ) + + vocab_size = min(logits_shape[1], bitmask_shape[1] * BITS_PER_BLOCK) + + num_rows = None + if isinstance(indices, list) or isinstance(indices, torch.Tensor): + indices = torch.tensor(indices, dtype=torch.int32, device=logits.device) + num_rows = indices.shape[0] + else: + assert ( + logits_shape[0] == bitmask_shape[0] + ), f"batch size mismatch: logits {logits_shape[0]} vs bitmask {bitmask_shape[0]}" + num_rows = logits_shape[0] + + if NUM_SMS > 0: + grid = (NUM_SMS,) + else: + num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE) + grid = (num_rows * num_blocks,) + NUM_SMS = triton.next_power_of_2(grid[0]) + + apply_token_bitmask_inplace_kernel[grid]( + logits, + bitmask, + indices, + num_rows, + vocab_size, + logits_shape[1], + bitmask_shape[1], + NUM_SMS, + BLOCK_SIZE, + num_warps=BLOCK_SIZE // 32 // (16 // logits.element_size()), + num_stages=3, + ) diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py new file mode 100644 index 000000000..7b101df4f --- /dev/null +++ b/python/sglang/srt/constrained/xgrammar_backend.py @@ -0,0 +1,239 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Constrained decoding with xgrammar backend.""" + +import json +import logging +from typing import List, Optional, Tuple, Union + +import torch +from xgrammar import ( + CompiledGrammar, + GrammarCompiler, + GrammarMatcher, + StructuralTagItem, + TokenizerInfo, + allocate_token_bitmask, +) + +from sglang.srt.constrained.base_grammar_backend import ( + INVALID_GRAMMAR_OBJ, + BaseGrammarBackend, + BaseGrammarObject, +) +from sglang.srt.utils import is_hip + +_is_hip = is_hip() +if _is_hip: + from sgl_kernel import apply_token_bitmask_inplace_cuda +else: + from sglang.srt.constrained.triton_ops.bitmask_ops import ( + apply_token_bitmask_inplace_triton, + ) +logger = logging.getLogger(__name__) + + +MAX_ROLLBACK_TOKENS = 200 + + +class XGrammarGrammar(BaseGrammarObject): + + def __init__( + self, + matcher: GrammarMatcher, + vocab_size: int, + ctx: CompiledGrammar, + override_stop_tokens: Optional[Union[List[int], int]], + key_string: Optional[str] = None, # TODO (sk): for debugging, remove later + ) -> None: + self.matcher = matcher + self.vocab_size = vocab_size + self.ctx = ctx + self.override_stop_tokens = override_stop_tokens + self.finished = False + self.accepted_tokens = [] + self.key_string = key_string + + def accept_token(self, token: int): + if not self.is_terminated(): + accepted = self.matcher.accept_token(token) + if not accepted: + # log for debugging + raise ValueError( + f"Tokens not accepted: {token}\n" + f"Accepted tokens: {self.accepted_tokens}\n" + f"Key string: {self.key_string}" + ) + else: + self.accepted_tokens.append(token) + + def rollback(self, k: int): + self.matcher.rollback(k) + self.accepted_tokens = self.accepted_tokens[:-k] + + def is_terminated(self): + return self.matcher.is_terminated() + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device + ) -> torch.Tensor: + return allocate_token_bitmask(batch_size, vocab_size) + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + self.matcher.fill_next_token_bitmask(vocab_mask, idx) + + @staticmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor: + return vocab_mask.to(device, non_blocking=True) + + def apply_vocab_mask(self, logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + if logits.device.type == "cuda": + if _is_hip: + apply_token_bitmask_inplace_cuda(logits, vocab_mask) + else: + apply_token_bitmask_inplace_triton(logits, vocab_mask) + elif logits.device.type == "cpu" and self.apply_vocab_mask_cpu: + self.apply_vocab_mask_cpu(logits, vocab_mask) + else: + raise RuntimeError(f"Unsupported device: {logits.device.type}") + + def copy(self): + matcher = GrammarMatcher( + self.ctx, + max_rollback_tokens=MAX_ROLLBACK_TOKENS, + override_stop_tokens=self.override_stop_tokens, + ) + return XGrammarGrammar( + matcher, + self.vocab_size, + self.ctx, + self.override_stop_tokens, + self.key_string, + ) + + def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]: + s = self.matcher.find_jump_forward_string() + if s: + return [], s + return None + + def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]: + _, data = helper + return data, -1 + + def jump_and_retokenize( + self, old_output_ids: List[int], new_output_ids: List[int], next_state: int + ): + k = 0 + for i, old_id in enumerate(old_output_ids): + if old_id == new_output_ids[i]: + k = i + 1 + else: + break + + # rollback to the last token that is the same + if k < len(old_output_ids): + self.matcher.rollback(len(old_output_ids) - k) + + for i in range(k, len(new_output_ids)): + assert self.matcher.accept_token(new_output_ids[i]) + + def __repr__(self): + return f"XGrammarGrammar({self.key_string=}, {self.accepted_tokens=})" + + +class XGrammarGrammarBackend(BaseGrammarBackend): + def __init__( + self, + tokenizer, + vocab_size: int, + model_eos_token_ids: Optional[List[int]] = None, + ): + super().__init__() + + if hasattr(tokenizer, "init_xgrammar"): + # For special tokenizer + tokenizer_info, override_stop_tokens = tokenizer.init_xgrammar() + else: + # Create TokenizerInfo with model's EOS tokens as the authoritative stop tokens + # This ensures consistency between what the model considers EOS and what XGrammar uses + tokenizer_info = TokenizerInfo.from_huggingface( + tokenizer, vocab_size=vocab_size, stop_token_ids=model_eos_token_ids + ) + override_stop_tokens = None + + self.grammar_compiler = GrammarCompiler(tokenizer_info=tokenizer_info) + self.vocab_size = vocab_size + self.override_stop_tokens = override_stop_tokens + + def _from_context(self, ctx: CompiledGrammar, key_string: str) -> XGrammarGrammar: + matcher = GrammarMatcher( + ctx, + max_rollback_tokens=MAX_ROLLBACK_TOKENS, + override_stop_tokens=self.override_stop_tokens, + ) + return XGrammarGrammar( + matcher, self.vocab_size, ctx, self.override_stop_tokens, key_string + ) + + def dispatch_json(self, key_string: str) -> Optional[XGrammarGrammar]: + try: + if key_string == "$$ANY$$": + # Note: This builtin JSON grammar includes *all* valid JSON (including, for example, arrays at the root) + ctx = self.grammar_compiler.compile_builtin_json_grammar() + else: + ctx = self.grammar_compiler.compile_json_schema(schema=key_string) + + except (RuntimeError, json.decoder.JSONDecodeError) as e: + logging.error(f"Hit invalid json_schema: {key_string=}, {e=}") + return INVALID_GRAMMAR_OBJ + return self._from_context(ctx, key_string) + + def dispatch_ebnf(self, key_string: str) -> Optional[XGrammarGrammar]: + try: + ctx = self.grammar_compiler.compile_grammar(key_string) + except RuntimeError as e: + logging.error(f"Hit invalid ebnf: {key_string=}, {e=}") + return INVALID_GRAMMAR_OBJ + return self._from_context(ctx, key_string) + + def dispatch_regex(self, key_string: str) -> Optional[XGrammarGrammar]: + try: + ctx = self.grammar_compiler.compile_regex(key_string) + except RuntimeError as e: + logging.error(f"Hit invalid regex: {key_string=}, {e=}") + return INVALID_GRAMMAR_OBJ + return self._from_context(ctx, key_string) + + def dispatch_structural_tag(self, key_string: str) -> Optional[XGrammarGrammar]: + try: + structural_tag = json.loads(key_string) + tags = [ + StructuralTagItem( + begin=structure["begin"], + schema=json.dumps(structure["schema"]), + end=structure["end"], + ) + for structure in structural_tag["structures"] + ] + ctx = self.grammar_compiler.compile_structural_tag( + tags, structural_tag["triggers"] + ) + except (RuntimeError, json.decoder.JSONDecodeError) as e: + logging.error(f"Hit invalid structural_tag: {key_string=}, {e=}") + return INVALID_GRAMMAR_OBJ + return self._from_context(ctx, key_string) + + def reset(self): + self.grammar_compiler.clear_cache() diff --git a/python/sglang/srt/custom_op.py b/python/sglang/srt/custom_op.py new file mode 100644 index 000000000..ea3c06e6d --- /dev/null +++ b/python/sglang/srt/custom_op.py @@ -0,0 +1,102 @@ +from torch import nn + +from sglang.srt.utils import ( + cpu_has_amx_support, + is_cpu, + is_cuda, + is_hip, + is_npu, + is_xpu, +) + +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_cpu = is_cpu() +_is_cpu_amx_available = cpu_has_amx_support() +_is_npu = is_npu() +_is_xpu = is_xpu() + + +class CustomOp(nn.Module): + def __init__(self): + super().__init__() + self._forward_method = self.dispatch_forward() + + # States for torch.compile + self._original_forward_method = None + self.is_torch_compile = False + + def enter_torch_compile(self, num_tokens: int): + # Skip if Op is already entered compile mode. + # NOTE(alcanderian): Some Ops(for example RotaryEmbedding) will be reused + # among layers and `enter_torch_compile` will be called many times. + # We should prevent `self._original_forward_method` from being overridden when + # it is not the first time `enter_torch_compile` called. + if self.is_torch_compile: + return + + self._original_forward_method = self._forward_method + # NOTE: Temporarily workaround MoE + # The performance of torch.compile on this layer is not always good when bs > 1, + # so we decide to only use torch.compile when bs=1 + if "FusedMoE" in self.__class__.__name__: + if num_tokens == 1: + from sglang.srt.layers.moe.fused_moe_native import ( + fused_moe_forward_native, + ) + + self._forward_method = fused_moe_forward_native + elif "TopK" in self.__class__.__name__: + if num_tokens == 1: + self._forward_method = self.forward_native + else: + self._forward_method = self.forward_native + self.is_torch_compile = True + + def leave_torch_compile(self): + # Skip if Op is already exited compile mode. + if not self.is_torch_compile: + return + + self._forward_method = self._original_forward_method + self._original_forward_method = None + self.is_torch_compile = False + + # Please do not override this method, because `self._forward_method` can change when in torch compile mode + def forward(self, *args, **kwargs): + return self._forward_method(*args, **kwargs) + + def forward_native(self, *args, **kwargs): + raise NotImplementedError + + def forward_cuda(self, *args, **kwargs): + raise NotImplementedError + + def forward_npu(self, *args, **kwargs): + raise NotImplementedError + + def forward_hip(self, *args, **kwargs): + return self.forward_cuda(*args, **kwargs) + + def forward_xpu(self, *args, **kwargs): + return self.forward_native(*args, **kwargs) + + def forward_hpu(self, *args, **kwargs): + return self.forward_native(*args, **kwargs) + + def forward_cpu(self, *args, **kwargs): + return self.forward_native(*args, **kwargs) + + def dispatch_forward(self): + if _is_cuda: + return self.forward_cuda + elif _is_hip: + return self.forward_hip + elif _is_cpu and _is_cpu_amx_available: + return self.forward_cpu + elif _is_npu: + return self.forward_npu + elif _is_xpu: + return self.forward_xpu + else: + return self.forward_native diff --git a/python/sglang/srt/debug_utils/__init__.py b/python/sglang/srt/debug_utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/debug_utils/dump_comparator.py b/python/sglang/srt/debug_utils/dump_comparator.py new file mode 100644 index 000000000..aca9c3b7a --- /dev/null +++ b/python/sglang/srt/debug_utils/dump_comparator.py @@ -0,0 +1,168 @@ +import argparse +import functools +from pathlib import Path + +import polars as pl +import torch + +from sglang.srt.debug_utils.dump_loader import find_row, read_meta +from sglang.srt.debug_utils.dumper import get_truncated_value + + +def main(args): + df_target = read_meta(args.target_path) + df_target = df_target.sort("rank", "dump_index") + df_target = df_target.filter( + (pl.col("forward_pass_id") >= args.start_id) + & (pl.col("forward_pass_id") <= args.end_id) + ) + assert all( + c in df_target.columns + for c in ["rank", "forward_pass_id", "dump_index", "name"] + ) + + df_baseline = read_meta(args.baseline_path) + print("df_target", df_target) + print("df_baseline", df_baseline) + + for row in df_target.iter_rows(named=True): + path_target = Path(args.target_path) / row["filename"] + + row_baseline = find_row( + df_baseline, + conditions=dict( + forward_pass_id=row["forward_pass_id"] + - args.start_id + + args.baseline_start_id, + **{ + k: v + for k, v in row.items() + if k not in ["forward_pass_id", "dump_index", "filename"] + }, + ), + ) + + if row_baseline is None: + print(f"Skip: target={str(path_target)} since no baseline") + x_target = _load_object(path_target) + if x_target is not None: + print(f"x_target(sample)={get_truncated_value(x_target)}") + continue + + path_baseline = Path(args.baseline_path) / row_baseline["filename"] + print(f"Check: target={str(path_target)} baseline={str(path_baseline)}") + check_tensor_pair( + path_baseline=path_baseline, path_target=path_target, name=row["name"] + ) + print() + + +def check_tensor_pair(path_baseline, path_target, name=""): + x_baseline = _load_object(path_baseline) + x_target = _load_object(path_target) + + print( + f"Raw " + f"[shape] {x_baseline.shape} vs {x_target.shape}\t" + f"[dtype] {x_baseline.dtype} vs {x_target.dtype}" + ) + + x_baseline, x_target = _comparison_preprocessor(x_baseline, x_target, name=name) + x_baseline = _try_unify_shape(x_baseline, target_shape=x_target.shape) + + print( + f"After preprocessor " + f"[shape] {x_baseline.shape} vs {x_target.shape}\t" + f"[dtype] {x_baseline.dtype} vs {x_target.dtype}" + ) + + x_target = x_target.float() + x_baseline = x_baseline.float() + + for name, fn in ( + ("mean", torch.mean), + ("std", torch.std), + ("min", torch.min), + ("max", torch.max), + ("p1", functools.partial(torch.quantile, q=0.01)), + ("p5", functools.partial(torch.quantile, q=0.05)), + ("p95", functools.partial(torch.quantile, q=0.95)), + ("p99", functools.partial(torch.quantile, q=0.99)), + ): + value_baseline = fn(x_baseline).item() + value_target = fn(x_target).item() + print( + f"[{name}] {value_baseline :.4f} vs {value_target:.4f} (diff: {value_target - value_baseline:.4f})" + ) + + if x_baseline.shape != x_target.shape: + print(f"⚠️ Shape mismatch") + return + + raw_abs_diff = (x_target - x_baseline).abs() + + max_abs_diff = raw_abs_diff.max().item() + mean_abs_diff = raw_abs_diff.mean().item() + rel_diff = _calc_rel_diff(x_target, x_baseline) + + needs_print = max_abs_diff > 1e-3 + + print( + "\t".join( + f"{'❌' if value > 1e-3 else '✅'} {name}={value}" + for name, value in [ + ("rel_diff", rel_diff), + ("max_abs_diff", max_abs_diff), + ("mean_abs_diff", mean_abs_diff), + ] + ) + ) + + if needs_print: + print(f"x_baseline(sample)={get_truncated_value(x_baseline)}") + print(f"x_target(sample)={get_truncated_value(x_target)}") + + +def _try_unify_shape(x: torch.Tensor, target_shape): + x_shape = x.shape + num_dim_to_remove = len(x_shape) - len(target_shape) + if (x_shape[num_dim_to_remove:] == target_shape) and all( + val == 1 for val in x_shape[:num_dim_to_remove] + ): + out = functools.reduce(lambda a, _: a.squeeze(0), range(num_dim_to_remove), x) + print(f"Unify shape: {x_shape} -> {out.shape} (to match {target_shape})") + return out + + return x + + +# Copied from DeepGEMM +def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor): + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + sim = 2 * (x * y).sum() / denominator + return 1 - sim + + +def _comparison_preprocessor(x_baseline, x_target, name): + # can insert arbitrary adhoc postprocessing logic here + return x_baseline, x_target + + +def _load_object(path): + x = torch.load(path, weights_only=False) + if not isinstance(x, torch.Tensor): + print(f"Skip load {path} since {type(x)=} is not a Tensor") + return None + return x.cuda() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--baseline-path", type=str) + parser.add_argument("--target-path", type=str) + parser.add_argument("--start-id", type=int, default=0) + parser.add_argument("--end-id", type=int, default=1000000) + parser.add_argument("--baseline-start-id", type=int, default=0) + args = parser.parse_args() + main(args) diff --git a/python/sglang/srt/debug_utils/dump_loader.py b/python/sglang/srt/debug_utils/dump_loader.py new file mode 100644 index 000000000..8e6f2c79b --- /dev/null +++ b/python/sglang/srt/debug_utils/dump_loader.py @@ -0,0 +1,97 @@ +import functools +import os +from pathlib import Path +from typing import Any, Dict + +import polars as pl +import torch + + +class DumpLoader: + def __init__(self): + directory = os.environ.get("SGLANG_DUMP_LOADER_DIR") + + self._enable = directory is not None + if self._enable: + self._directory = Path(directory) + self._df = read_meta(directory) + + @property + def enable(self): + return self._enable + + def load(self, name, **kwargs): + assert self._enable, "Please call DumpLoader.load only when it is enabled" + + from sglang.srt.debug_utils.dumper import dumper + + forward_pass_id = dumper._forward_pass_id + conditions = dict(name=name, forward_pass_id=forward_pass_id, **kwargs) + row = find_row(self._df, conditions=conditions) + assert ( + row is not None + ), f"DumpLoader cannot find row given query {name=} {kwargs=} {self._directory=}" + + path = self._directory / row["filename"] + output = torch.load(path, weights_only=False) + + print( + f"[DumpLoader] load from {path=} (query: {name=} {kwargs=}, output: {type(output)})" + ) + return output + + +def read_meta(directory): + directory = Path(directory) + assert directory.is_dir(), f"{directory=} should be a directory" + + rows = [] + for p in directory.glob("*.pt"): + full_kwargs = {} + for kv in p.stem.split("___"): + k, v = kv.split("=") + full_kwargs[k] = v + rows.append( + { + "filename": str(p.name), + **full_kwargs, + } + ) + + df = pl.DataFrame(rows) + df = df.with_columns( + pl.col("forward_pass_id").cast(int), + pl.col("rank").cast(int), + pl.col("dump_index").cast(int), + ) + return df + + +def find_row(df, conditions: Dict[str, Any]): + df_sub = df.filter( + functools.reduce( + lambda a, b: a & b, + [ + pl.col(col) == _cast_to_polars_dtype(conditions[col], df.schema[col]) + for col in conditions.keys() + ], + ) + ) + assert len(df_sub) <= 1 + return df_sub.to_dicts()[0] if len(df_sub) > 0 else None + + +def _cast_to_polars_dtype(value, target_dtype): + if target_dtype in (pl.Int64, pl.Int32, pl.UInt64, pl.UInt32): + return int(value) + elif target_dtype in (pl.Float64, pl.Float32): + return float(value) + elif target_dtype == pl.Boolean: + return bool(value) + elif target_dtype == pl.String: + return str(value) + else: + return value + + +dump_loader = DumpLoader() diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py new file mode 100644 index 000000000..8a9808bb7 --- /dev/null +++ b/python/sglang/srt/debug_utils/dumper.py @@ -0,0 +1,116 @@ +import os +import time +from pathlib import Path +from typing import Optional + +import torch +import torch.distributed as dist + + +class _Dumper: + """Utility to dump tensors, which can be useful when comparison checking models. + + Example usage: + dumper.on_forward_pass_start() + dumper.dump("layer_start__hidden_states", hidden_states, layer_id=self.layer_id) + + Import from non-SGLang system: + ``` + import sys + sys.path.append("/YOUR_PATH/sglang/python/sglang/srt/debug_utils") + from dumper import dumper + ``` + + Related: `sglang.srt.debug_utils.dump_comparator` for dump comparison + """ + + def __init__(self): + # Do not import `sglang` to make this file standalone + self._enable = bool(int(os.environ.get("SGLANG_DUMPER_ENABLE", "1"))) + self._base_dir = Path(os.environ.get("SGLANG_DUMPER_DIR", "/tmp")) + self._enable_write_file = bool( + int(os.environ.get("SGLANG_DUMPER_WRITE_FILE", "1")) + ) + self._partial_name: Optional[str] = None + self._dump_index = 0 + self._forward_pass_id = 0 + + def on_forward_pass_start(self): + self._forward_pass_id += 1 + print( + f"[Dumper] [{time.time()}] on_forward_pass_start id={self._forward_pass_id}" + ) + + def dump(self, name, value, **kwargs): + if not self._enable: + return + + assert ( + self._forward_pass_id >= 1 + ), "Do you forget to call `dumper.on_forward_pass_start()`?" + self._dump_index += 1 + + if self._partial_name is None: + self._partial_name = _get_partial_name() + + rank = _get_rank() + full_kwargs = dict( + forward_pass_id=self._forward_pass_id, + rank=rank, + name=name, + dump_index=self._dump_index, + **kwargs, + ) + full_filename = "___".join(f"{k}={v}" for k, v in full_kwargs.items()) + ".pt" + path = self._base_dir / f"sglang_dump_{self._partial_name}" / full_filename + + sample_value = get_truncated_value(value) + + print( + f"[Dumper] [{rank}, {time.time()}] {path} " + f"type={type(value)} " + f"shape={value.shape if isinstance(value, torch.Tensor) else None} " + f"dtype={value.dtype if isinstance(value, torch.Tensor) else None} " + f"sample_value={sample_value}" + ) + + if self._enable_write_file: + path.parent.mkdir(parents=True, exist_ok=True) + torch.save(value, str(path)) + + +def _get_partial_name(): + rank = _get_rank() + object_list = [str(time.time()) if rank == 0 else None] + if dist.is_initialized(): + dist.broadcast_object_list(object_list, device="cuda") + return object_list[0] + + +def _get_rank(): + if dist.is_initialized(): + return dist.get_rank() + else: + return 0 + + +def get_truncated_value(value): + if value is None: + return None + + if isinstance(value, tuple): + return [get_truncated_value(x) for x in value] + + if not isinstance(value, torch.Tensor): + return None + + if value.numel() < 200: + return value + + slices = [ + slice(0, 5) if dim_size > 200 else slice(None) for dim_size in value.shape + ] + return value[tuple(slices)] + + +dumper = _Dumper() diff --git a/python/sglang/srt/debug_utils/text_comparator.py b/python/sglang/srt/debug_utils/text_comparator.py new file mode 100644 index 000000000..3a6df19b9 --- /dev/null +++ b/python/sglang/srt/debug_utils/text_comparator.py @@ -0,0 +1,234 @@ +import argparse +import hashlib +import json +from pathlib import Path + +import polars as pl + +_DESCRIPTION = """Compare and find differences to benchmark outputs. + +Supported inputs: +* The samples jsonl from `lm_eval --log_samples --output_path FOLDER_NAME` +* The output from `gsm8k/bench_sglang.py --raw-result-file FILE_NAME` (or mmlu) +""" + + +def main(args): + if args.data_type == "simple_evals": + df_input = _compute_df_input_mode_simple_evals(args) + else: + df_input = _transform_df_input(_compute_df_raw(args)) + + assert all( + c in df_input.columns + for c in ["category", "trial_index", "prompt_id", "prompt", "output", "correct"] + ) + + df_meta = _compute_df_meta(df_input) + + df_correctness_per_trial = df_input.group_by( + "category", "trial_index", maintain_order=True + ).agg(pl.col("correct").mean()) + df_correctness_delta = ( + df_meta.group_by("correctness_delta").len().sort("correctness_delta") + ) + df_good_to_bad = df_meta.filter(pl.col("correctness_delta") < 0) + df_bad_to_good = df_meta.filter(pl.col("correctness_delta") > 0) + + print(f"Dump output to {args.output_path}") + Path(args.output_path).write_text( + json.dumps( + dict( + df_meta=df_meta.to_dicts(), + df_good_to_bad=df_good_to_bad.to_dicts(), + df_bad_to_good=df_bad_to_good.to_dicts(), + ), + indent=4, + ), + ) + + if not args.disable_print_details: + with pl.Config( + fmt_str_lengths=10000, + tbl_cols=-1, + tbl_rows=-1, + tbl_width_chars=-1, + tbl_formatting="UTF8_FULL", + ): + print("====== Correctness per trial ======") + print(df_correctness_per_trial) + + print( + "====== Correctness Delta (-1.0 means all-right becomes all-wrong) ======" + ) + print(df_correctness_delta) + + for name, df in [ + ("Good->Bad", df_good_to_bad), + ("Bad->Good", df_bad_to_good), + ]: + print(f"====== Concrete Examples: {name} ======") + print(df) + + +def _compute_df_input_mode_simple_evals(args): + return pl.concat( + [ + _compute_df_input_one_mode_simple_evals(**info) + for info in _get_file_infos(args=args) + ] + ) + + +def _compute_df_input_one_mode_simple_evals(path, category, trial_index): + data = json.loads(Path(path).read_text()) + rows = [] + + for single_eval_result in data["metadata"]["single_eval_results"]: + prompt = single_eval_result["example_level_metadata"][ + "actual_queried_prompt_messages" + ] + score = single_eval_result["score"] + assert score in {0.0, 1.0}, f"{score=}" + + row = dict( + category=category, + trial_index=trial_index, + prompt_id=_compute_id_from_object(prompt), + prompt=json.dumps(prompt), + output=single_eval_result["example_level_metadata"]["response_text"], + correct=score == 1.0, + ) + rows.append(row) + + return pl.DataFrame(rows) + + +def _compute_id_from_object(obj): + if isinstance(obj, pl.Series): + obj = obj.to_list() + json_str = json.dumps(obj, sort_keys=True, ensure_ascii=False) + return hashlib.sha256(json_str.encode("utf-8")).hexdigest() + + +def _compute_df_raw(args): + return pl.concat( + [ + _read_df_raw( + path=info["path"], + category=info["category"], + trial_index=info["trial_index"], + ) + for info in _get_file_infos(args=args) + ] + ) + + +def _get_file_infos(args): + return [ + dict(path=path, category=category, trial_index=trial_index) + for category, paths in [ + ("baseline", args.baseline_path), + ("target", args.target_path), + ] + for trial_index, path in enumerate(paths) + ] + + +def _read_df_raw(path: str, category: str, trial_index: int): + return pl.read_ndjson(path).with_columns( + category=pl.lit(category), trial_index=trial_index + ) + + +def _transform_df_input(df: pl.DataFrame): + if "doc_id" in df.columns: + print("Transform mode: lm_eval") + + filter_names = df["filter"].unique(maintain_order=True).to_list() + if len(filter_names) > 1: + filter_name = filter_names[0] + print(f"Choose {filter_name=} among {filter_names}") + df = df.filter(pl.col("filter") == filter_name) + + df = df.select( + pl.col("category"), + pl.col("trial_index"), + prompt_id=pl.col("doc_id"), + prompt=pl.col("arguments").struct.field("gen_args_0").struct.field("arg_0"), + output=pl.col("resps").list.get(0).list.get(0), + correct=pl.col("exact_match").cast(bool), + ) + + return df + elif "prompt_id" in df.columns: + print("Transform mode: SGLang bench") + return df + else: + raise Exception( + f"Unknown data: {df.columns}. You may need to set `--data-type` if using e.g. simple_evals." + ) + + +def _compute_df_meta(df_input: pl.DataFrame): + df_input = df_input.sort("prompt_id", "category", "trial_index") + df_meta = pl.DataFrame( + [ + _handle_one_prompt(df_one_prompt) + for df_one_prompt in df_input.partition_by("prompt_id", maintain_order=True) + ] + ) + df_meta = df_meta.with_columns( + correctness_delta=pl.col("correctness_target") - pl.col("correctness_baseline"), + ) + df_meta = df_meta.sort("correctness_delta", "output_same_prefix_len") + return df_meta + + +def _handle_one_prompt(df_one_prompt: pl.DataFrame): + assert ( + len(set(_compute_id_from_object(obj) for obj in df_one_prompt["prompt"])) == 1 + ) + + df_baseline = df_one_prompt.filter(pl.col("category") == "baseline") + df_target = df_one_prompt.filter(pl.col("category") == "target") + + outputs_baseline = df_baseline["output"].to_list() + outputs_target = df_target["output"].to_list() + + output_same_prefix_len = max( + _compute_str_prefix_len(output_baseline, output_target) + for output_baseline in outputs_baseline + for output_target in outputs_target + ) + + return dict( + prompt_id=df_one_prompt[0, "prompt_id"], + correctness_baseline=df_baseline["correct"].mean(), + correctness_target=df_target["correct"].mean(), + output_same_prefix_len=output_same_prefix_len, + prompt=df_one_prompt[0, "prompt"], + outputs_baseline=outputs_baseline, + outputs_target=outputs_target, + ) + + +def _compute_str_prefix_len(a: str, b: str) -> int: + min_len = min(len(a), len(b)) + for i in range(min_len): + if a[i] != b[i]: + return i + return min_len + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=_DESCRIPTION) + parser.add_argument("--data-type", type=str, default="auto") + parser.add_argument("--baseline-path", type=str, nargs="+") + parser.add_argument("--target-path", type=str, nargs="+") + parser.add_argument( + "--output-path", type=str, default="/tmp/text_comparator_output.json" + ) + parser.add_argument("--disable-print-details", action="store_true") + args = parser.parse_args() + main(args) diff --git a/python/sglang/srt/disaggregation/ascend/__init__.py b/python/sglang/srt/disaggregation/ascend/__init__.py new file mode 100644 index 000000000..2550f91a4 --- /dev/null +++ b/python/sglang/srt/disaggregation/ascend/__init__.py @@ -0,0 +1,6 @@ +from sglang.srt.disaggregation.ascend.conn import ( + AscendKVBootstrapServer, + AscendKVManager, + AscendKVReceiver, + AscendKVSender, +) diff --git a/python/sglang/srt/disaggregation/ascend/conn.py b/python/sglang/srt/disaggregation/ascend/conn.py new file mode 100644 index 000000000..b0009fc7c --- /dev/null +++ b/python/sglang/srt/disaggregation/ascend/conn.py @@ -0,0 +1,117 @@ +import concurrent.futures +import logging +from typing import List, Tuple + +import numpy as np +import numpy.typing as npt + +from sglang.srt.disaggregation.ascend.transfer_engine import AscendTransferEngine +from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous +from sglang.srt.disaggregation.mooncake.conn import ( + MooncakeKVBootstrapServer, + MooncakeKVManager, + MooncakeKVReceiver, + MooncakeKVSender, +) +from sglang.srt.utils import get_local_ip_by_remote + +logger = logging.getLogger(__name__) + + +class AscendKVManager(MooncakeKVManager): + def init_engine(self): + # TransferEngine initialized on ascend. + local_ip = get_local_ip_by_remote() + self.engine = AscendTransferEngine( + hostname=local_ip, + npu_id=self.kv_args.gpu_id, + disaggregation_mode=self.disaggregation_mode, + ) + + def register_buffer_to_engine(self): + self.engine.batch_register(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens) + # The Ascend backend optimize batch registration for small memory blocks. + self.engine.batch_register( + self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens + ) + + def send_kvcache( + self, + mooncake_session_id: str, + prefill_kv_indices: npt.NDArray[np.int32], + dst_kv_ptrs: list[int], + dst_kv_indices: npt.NDArray[np.int32], + executor: concurrent.futures.ThreadPoolExecutor, + ): + # Group by indices + prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( + prefill_kv_indices, dst_kv_indices + ) + + num_layers = len(self.kv_args.kv_data_ptrs) + layers_params = [ + ( + self.kv_args.kv_data_ptrs[layer_id], + dst_kv_ptrs[layer_id], + self.kv_args.kv_item_lens[layer_id], + ) + for layer_id in range(num_layers) + ] + + def set_transfer_blocks( + src_ptr: int, dst_ptr: int, item_len: int + ) -> List[Tuple[int, int, int]]: + transfer_blocks = [] + for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks): + src_addr = src_ptr + int(prefill_index[0]) * item_len + dst_addr = dst_ptr + int(decode_index[0]) * item_len + length = item_len * len(prefill_index) + transfer_blocks.append((src_addr, dst_addr, length)) + return transfer_blocks + + # Worker function for processing a single layer + def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int: + transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len) + return self._transfer_data(mooncake_session_id, transfer_blocks) + + # Worker function for processing all layers in a batch + def process_layers(layers_params: List[Tuple[int, int, int]]) -> int: + transfer_blocks = [] + for src_ptr, dst_ptr, item_len in layers_params: + transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len)) + return self._transfer_data(mooncake_session_id, transfer_blocks) + + if self.enable_custom_mem_pool: + futures = [ + executor.submit( + process_layer, + src_ptr, + dst_ptr, + item_len, + ) + for (src_ptr, dst_ptr, item_len) in layers_params + ] + for future in concurrent.futures.as_completed(futures): + status = future.result() + if status != 0: + for f in futures: + f.cancel() + return status + else: + # Combining all layers' params in one batch transfer is more efficient + # compared to using multiple threads + return process_layers(layers_params) + + return 0 + + +class AscendKVSender(MooncakeKVSender): + pass + + +class AscendKVReceiver(MooncakeKVReceiver): + pass + + +class AscendKVBootstrapServer(MooncakeKVBootstrapServer): + pass diff --git a/python/sglang/srt/disaggregation/ascend/transfer_engine.py b/python/sglang/srt/disaggregation/ascend/transfer_engine.py new file mode 100644 index 000000000..0ccffffd6 --- /dev/null +++ b/python/sglang/srt/disaggregation/ascend/transfer_engine.py @@ -0,0 +1,58 @@ +import logging +import os +from typing import List, Optional + +from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine +from sglang.srt.disaggregation.utils import DisaggregationMode + +logger = logging.getLogger(__name__) + + +class AscendTransferEngine(MooncakeTransferEngine): + + def __init__( + self, hostname: str, npu_id: int, disaggregation_mode: DisaggregationMode + ): + try: + from mf_adapter import TransferEngine + except ImportError as e: + raise ImportError( + "Please install mf_adapter, for details, see docs/backend/pd_disaggregation.md" + ) from e + + self.engine = TransferEngine() + self.hostname = hostname + self.npu_id = npu_id + + # Centralized storage address of the AscendTransferEngine + self.store_url = os.getenv("ASCEND_MF_STORE_URL") + if disaggregation_mode == DisaggregationMode.PREFILL: + self.role = "Prefill" + elif disaggregation_mode == DisaggregationMode.DECODE: + self.role = "Decode" + else: + logger.error(f"Unsupported DisaggregationMode: {disaggregation_mode}") + raise ValueError(f"Unsupported DisaggregationMode: {disaggregation_mode}") + self.session_id = f"{self.hostname}:{self.engine.get_rpc_port()}" + self.initialize() + + def initialize(self) -> None: + """Initialize the ascend transfer instance.""" + ret_value = self.engine.initialize( + self.store_url, + self.session_id, + self.role, + self.npu_id, + ) + if ret_value != 0: + logger.error("Ascend Transfer Engine initialization failed.") + raise RuntimeError("Ascend Transfer Engine initialization failed.") + + def batch_register(self, ptrs: List[int], lengths: List[int]): + try: + ret_value = self.engine.batch_register_memory(ptrs, lengths) + except Exception: + # Mark register as failed + ret_value = -1 + if ret_value != 0: + logger.debug(f"Ascend memory registration for ptr {ptrs} failed.") diff --git a/python/sglang/srt/disaggregation/base/__init__.py b/python/sglang/srt/disaggregation/base/__init__.py new file mode 100644 index 000000000..ef0f797fc --- /dev/null +++ b/python/sglang/srt/disaggregation/base/__init__.py @@ -0,0 +1,8 @@ +from sglang.srt.disaggregation.base.conn import ( + BaseKVBootstrapServer, + BaseKVManager, + BaseKVReceiver, + BaseKVSender, + KVArgs, + KVPoll, +) diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py new file mode 100644 index 000000000..3f5877ea3 --- /dev/null +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, List, Optional + +import numpy as np +import numpy.typing as npt + +from sglang.srt.server_args import ServerArgs + +if TYPE_CHECKING: + from sglang.srt.disaggregation.utils import DisaggregationMode + + +class KVArgs: + engine_rank: int + kv_data_ptrs: List[int] + kv_data_lens: List[int] + kv_item_lens: List[int] + aux_data_ptrs: List[int] + aux_data_lens: List[int] + aux_item_lens: List[int] + ib_device: str + ib_traffic_class: str + gpu_id: int + # for different tp + decode_tp_size: int + kv_head_num: int + page_size: int + # for pp prefill + prefill_pp_size: int + pp_rank: int + prefill_start_layer: int + # for system dp + system_dp_rank: int + + +class KVPoll: + Failed = 0 + Bootstrapping = 1 + WaitingForInput = 2 + Transferring = 3 + Success = 4 + + +class BaseKVManager(ABC): + """Base class for managing transfers states""" + + @abstractmethod + def __init__( + self, + args: KVArgs, + disaggregation_mode: DisaggregationMode, + server_args: ServerArgs, + is_mla_backend: Optional[bool] = False, + ): ... + + +class BaseKVSender(ABC): + + @abstractmethod + def __init__( + self, + mgr: BaseKVManager, + bootstrap_addr: str, + bootstrap_room: int, + dest_tp_ranks: List[int], + pp_rank: int, + ): ... + + @abstractmethod + def init(self, num_kv_indices: int, aux_index: Optional[int] = None): + """ + Notify the decoder server about the kv indices length and aux index + """ + ... + + @abstractmethod + def send(self, kv_indices: npt.NDArray[np.int32]): + """ + Send the kv cache at the given kv indices to the decoder server + """ + ... + + @abstractmethod + def poll(self) -> KVPoll: + """ + Check the status of the kv cache transfer + """ + ... + + @abstractmethod + def failure_exception(self): + """ + Raise an exception if the kv cache transfer fails + """ + ... + + +class BaseKVReceiver(ABC): + + @abstractmethod + def __init__( + self, + mgr: BaseKVManager, + bootstrap_addr: str, + bootstrap_room: Optional[int] = None, + ): ... + + @abstractmethod + def init(self, kv_indices: npt.NDArray[np.int32], aux_index: Optional[int] = None): + """ + Notify the prefill server about the kv indices and aux index + """ + ... + + @abstractmethod + def poll(self) -> KVPoll: + """ + Check the status of the kv cache transfer + """ + ... + + @abstractmethod + def failure_exception(self): + """ + Raise an exception if the kv cache transfer fails + """ + ... + + +class BaseKVBootstrapServer(ABC): + @abstractmethod + def __init__(self, host: str, port: int): ... diff --git a/python/sglang/srt/disaggregation/common/__init__.py b/python/sglang/srt/disaggregation/common/__init__.py new file mode 100644 index 000000000..8294c3892 --- /dev/null +++ b/python/sglang/srt/disaggregation/common/__init__.py @@ -0,0 +1,5 @@ +from sglang.srt.disaggregation.common.conn import ( + CommonKVBootstrapServer, + CommonKVManager, + CommonKVReceiver, +) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py new file mode 100644 index 000000000..10b6093b9 --- /dev/null +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -0,0 +1,438 @@ +from __future__ import annotations + +import asyncio +import logging +import socket +import threading +from functools import cache +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import numpy.typing as npt +import requests +import zmq +from aiohttp import web + +from sglang.srt.disaggregation.base.conn import ( + BaseKVBootstrapServer, + BaseKVManager, + BaseKVReceiver, + BaseKVSender, + KVArgs, + KVPoll, +) +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import ( + format_tcp_address, + get_free_port, + get_ip, + get_local_ip_by_remote, + is_valid_ipv6_address, + maybe_wrap_ipv6_address, +) + +logger = logging.getLogger(__name__) + + +class CommonKVManager(BaseKVManager): + def __init__( + self, + args: KVArgs, + disaggregation_mode: DisaggregationMode, + server_args: ServerArgs, + is_mla_backend: Optional[bool] = False, + ): + self.kv_args = args + self.is_mla_backend = is_mla_backend + self.disaggregation_mode = disaggregation_mode + # for p/d multi node infer + self.bootstrap_host = server_args.host + self.bootstrap_port = server_args.disaggregation_bootstrap_port + self.dist_init_addr = server_args.dist_init_addr + self.tp_size = server_args.tp_size + self.dp_size = server_args.dp_size + self.enable_dp_attention = server_args.enable_dp_attention + if not server_args.enable_dp_attention and server_args.dp_size != 1: + raise ValueError( + "If dp_attention is not enabled, dp size must be 1 in disaggregation mode." + ) + + self.rank_port = get_free_port() + if self.disaggregation_mode == DisaggregationMode.PREFILL: + self._register_to_bootstrap() + elif self.disaggregation_mode == DisaggregationMode.DECODE: + self.connection_pool: Dict[str, Dict[str, Union[str, int]]] = {} + self.prefill_tp_size_table: Dict[str, int] = {} + self.prefill_dp_size_table: Dict[str, int] = {} + else: + raise ValueError( + f"Unsupported DisaggregationMode: {self.disaggregation_mode}" + ) + + def _register_to_bootstrap(self): + """Register KVSender to bootstrap server via HTTP POST.""" + if self.dist_init_addr: + # multi node: bootstrap server's host is dist_init_addr + if self.dist_init_addr.startswith("["): # [ipv6]:port or [ipv6] + if self.dist_init_addr.endswith("]"): + host = self.dist_init_addr + else: + host, _ = self.dist_init_addr.rsplit(":", 1) + else: + host = socket.gethostbyname(self.dist_init_addr.rsplit(":", 1)[0]) + else: + # single node: bootstrap server's host is same as http server's host + host = self.bootstrap_host + host = maybe_wrap_ipv6_address(host) + + bootstrap_server_url = f"{host}:{self.bootstrap_port}" + url = f"http://{bootstrap_server_url}/route" + payload = { + "role": "Prefill", + "tp_size": self.tp_size, + "dp_size": self.dp_size, + "rank_ip": get_local_ip_by_remote(), + "rank_port": self.rank_port, + "engine_rank": self.kv_args.engine_rank, + } + + try: + response = requests.put(url, json=payload) + if response.status_code == 200: + logger.debug("Prefill successfully registered to bootstrap server.") + else: + logger.error( + f"Prefill Failed to connect to bootstrap server: {response.status_code}, {response.text}" + ) + except Exception as e: + logger.error(f"Prefill Failed to register to bootstrap server: {e}") + + @cache + def _connect(self, endpoint: str, is_ipv6: bool = False): + socket = zmq.Context().socket(zmq.PUSH) + if is_ipv6: + socket.setsockopt(zmq.IPV6, 1) + socket.connect(endpoint) + return socket + + +class CommonKVReceiver(BaseKVReceiver): + _ctx = zmq.Context() + _socket_cache = {} + _socket_locks = {} + _global_lock = threading.Lock() + + def __init__( + self, + mgr: BaseKVManager, + bootstrap_addr: str, + bootstrap_room: Optional[int] = None, + prefill_dp_rank: Optional[int] = None, + ): + self.bootstrap_room = bootstrap_room + self.bootstrap_addr = bootstrap_addr + self.kv_mgr = mgr + + if self.bootstrap_addr not in self.kv_mgr.prefill_dp_size_table: + self.prefill_tp_size, self.prefill_dp_size = ( + self._get_prefill_dp_size_from_server() + ) + if self.prefill_tp_size is None or self.prefill_dp_size is None: + logger.error( + f"Could not fetch prefill parallel info for bootstrap_addr: {self.bootstrap_addr}" + ) + else: + self.kv_mgr.prefill_tp_size_table[self.bootstrap_addr] = ( + self.prefill_tp_size + ) + self.kv_mgr.prefill_dp_size_table[self.bootstrap_addr] = ( + self.prefill_dp_size + ) + else: + self.prefill_tp_size = self.kv_mgr.prefill_tp_size_table[ + self.bootstrap_addr + ] + self.prefill_dp_size = self.kv_mgr.prefill_dp_size_table[ + self.bootstrap_addr + ] + + # Currently, we don't allow prefill instance and decode instance to + # have different TP sizes per DP rank, except for models using MLA. + local_tp_size_per_dp_rank = self.kv_mgr.tp_size // self.kv_mgr.dp_size + prefill_tp_size_per_dp_rank = self.prefill_tp_size // self.prefill_dp_size + if local_tp_size_per_dp_rank == prefill_tp_size_per_dp_rank: + self.target_tp_rank = ( + self.kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank + ) + self.required_dst_info_num = 1 + self.target_tp_ranks = [self.target_tp_rank] + elif local_tp_size_per_dp_rank > prefill_tp_size_per_dp_rank: + self.target_tp_rank = ( + self.kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank + ) // (local_tp_size_per_dp_rank // prefill_tp_size_per_dp_rank) + self.required_dst_info_num = ( + local_tp_size_per_dp_rank // prefill_tp_size_per_dp_rank + ) + self.target_tp_ranks = [self.target_tp_rank] + else: + assert ( + self.kv_mgr.is_mla_backend + ), "PD with different TP sizes per DP rank is not yet supported for non-MLA models" + + # For non-MLA models, one decode rank needs to retrieve KVCache from multiple prefill ranks for non MLA models; + self.target_tp_ranks = [ + rank + for rank in range( + (self.kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank) + * (prefill_tp_size_per_dp_rank // local_tp_size_per_dp_rank), + (self.kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank + 1) + * (prefill_tp_size_per_dp_rank // local_tp_size_per_dp_rank), + ) + ] + + # For MLA models, we can retrieve KVCache from only one prefill rank, but we still need to maintain + # multiple connections in the connection pool and have to send dummy requests to other prefill ranks, + # or the KVPoll will never be set correctly + self.target_tp_rank = self.target_tp_ranks[0] + self.required_dst_info_num = 1 + + if prefill_dp_rank is not None: + logger.debug(f"Targeting DP rank: {prefill_dp_rank}") + self.prefill_dp_rank = prefill_dp_rank + else: + self.prefill_dp_rank = bootstrap_room % self.prefill_dp_size + + # FIXME: alias here: target_dp_group -> prefill_dp_rank + self.target_dp_group = self.prefill_dp_rank + + # NOTE: key distinguished by bootstrap_addr, target_dp_group, and target_tp_rank + bootstrap_key = ( + f"{self.bootstrap_addr}_{self.target_dp_group}_{self.target_tp_rank}" + ) + + if bootstrap_key not in self.kv_mgr.connection_pool: + bootstrap_infos = [] + for target_tp_rank in self.target_tp_ranks: + bootstrap_info = self._get_bootstrap_info_from_server( + target_tp_rank, + self.target_dp_group, + ) + if bootstrap_info is not None: + # NOTE: only support MLA for now: select one prefill rank as real rank + bootstrap_info["is_dummy"] = not bool( + target_tp_rank == self.target_tp_rank + or self.target_tp_rank is None + ) + bootstrap_infos.append(bootstrap_info) + else: + logger.error( + f"Could not fetch bootstrap info for engine rank: {self.kv_mgr.kv_args.engine_rank} and target_dp_group: {self.target_dp_group}" + ) + self.bootstrap_infos = bootstrap_infos + + if len(self.bootstrap_infos) == 0: + logger.error( + f"Could not fetch bootstrap info for engine rank: {self.kv_mgr.kv_args.engine_rank}" + ) + else: + self.kv_mgr.connection_pool[bootstrap_key] = self.bootstrap_infos + # Register kv_args only once to prefill KVManager according to the info fetched from the bootstrap server + self._register_kv_args() + else: + self.bootstrap_infos = self.kv_mgr.connection_pool[bootstrap_key] + + assert len(self.bootstrap_infos) > 0 + + def _get_bootstrap_info_from_server(self, engine_rank, target_dp_group): + """Fetch the bootstrap info from the bootstrap server.""" + try: + url = f"http://{self.bootstrap_addr}/route?engine_rank={engine_rank}&target_dp_group={target_dp_group}" + response = requests.get(url) + if response.status_code == 200: + bootstrap_info = response.json() + return bootstrap_info + else: + logger.error( + f"Failed to get prefill server info: {response.status_code}, {response.text}" + ) + return None + except Exception as e: + logger.error(f"Error fetching prefill info from bootstrap: {e}") + return None + + def _get_prefill_dp_size_from_server(self) -> int: + """Fetch the prefill parallel info from the bootstrap server.""" + try: + url = f"http://{self.bootstrap_addr}/route?engine_rank={-1}&target_dp_group={-1}" + response = requests.get(url) + if response.status_code == 200: + prefill_parallel_info = response.json() + return int(prefill_parallel_info["prefill_tp_size"]), int( + prefill_parallel_info["prefill_dp_size"] + ) + else: + logger.error( + f"Failed to get prefill parallel info: {response.status_code}, {response.text}" + ) + return None + except Exception as e: + logger.error(f"Error fetching prefill parallel info from bootstrap: {e}") + return None + + @classmethod + def _connect(cls, endpoint: str, is_ipv6: bool = False): + with cls._global_lock: + if endpoint not in cls._socket_cache: + sock = cls._ctx.socket(zmq.PUSH) + if is_ipv6: + sock.setsockopt(zmq.IPV6, 1) + sock.connect(endpoint) + cls._socket_cache[endpoint] = sock + cls._socket_locks[endpoint] = threading.Lock() + return cls._socket_cache[endpoint], cls._socket_locks[endpoint] + + @classmethod + def _connect_to_bootstrap_server(cls, bootstrap_info: dict): + ip_address = bootstrap_info["rank_ip"] + port = bootstrap_info["rank_port"] + is_ipv6_address = is_valid_ipv6_address(ip_address) + sock, lock = cls._connect( + format_tcp_address(ip_address, port), is_ipv6=is_ipv6_address + ) + return sock, lock + + def _register_kv_args(self): + pass + + def failure_exception(self): + raise Exception("Fake KVReceiver Exception") + + +class CommonKVBootstrapServer(BaseKVBootstrapServer): + def __init__(self, host: str, port: int): + self.host = host + self.port = port + self.app = web.Application() + self.store = dict() + self.lock = asyncio.Lock() + self._setup_routes() + self.tp_size = None + self.dp_size = None + self.tp_size_per_dp_rank = None + self.prefill_port_table: Dict[int, Dict[int, Dict[str, Union[str, int]]]] = {} + + # Start bootstrap server + self.thread = threading.Thread(target=self._run_server, daemon=True) + self.run() + + def run(self): + self.thread.start() + + def _setup_routes(self): + self.app.router.add_route("*", "/route", self._handle_route) + + async def _handle_route(self, request: web.Request): + method = request.method + if method == "PUT": + return await self._handle_route_put(request) + elif method == "GET": + return await self._handle_route_get(request) + else: + return web.Response( + text="Method not allowed", status=405, content_type="application/json" + ) + + async def _handle_route_put(self, request: web.Request): + data = await request.json() + role = data["role"] + tp_size = data["tp_size"] + dp_size = data["dp_size"] + rank_ip = data["rank_ip"] + rank_port = int(data["rank_port"]) + engine_rank = int(data["engine_rank"]) + + if self.tp_size is None: + self.tp_size = tp_size + + if self.dp_size is None: + self.dp_size = dp_size + + tp_size_per_dp_rank = tp_size // dp_size + if self.tp_size_per_dp_rank == None: + self.tp_size_per_dp_rank = tp_size_per_dp_rank + + # Add lock to make sure thread-safe + if role == "Prefill": + dp_group = engine_rank // tp_size_per_dp_rank + tp_rank_in_dp_group = engine_rank % tp_size_per_dp_rank + + async with self.lock: + if dp_group not in self.prefill_port_table: + self.prefill_port_table[dp_group] = {} + + self.prefill_port_table[dp_group][tp_rank_in_dp_group] = { + "rank_ip": rank_ip, + "rank_port": rank_port, + } + logger.debug( + f"Register Prefill bootstrap: {engine_rank} with rank_ip: {rank_ip} and rank_port: {rank_port}" + ) + + return web.Response(text="OK", status=200) + + async def _handle_route_get(self, request: web.Request): + engine_rank = request.query.get("engine_rank") + target_dp_group = request.query.get("target_dp_group") + if not engine_rank or not target_dp_group: + return web.Response(text="Missing inputs for bootstrap server.", status=400) + + # Currently we use engine_rank == -1 and target_dp_group == -1 to sync dp size + if int(engine_rank) == -1 and int(target_dp_group) == -1: + prefill_parallel_info = { + "prefill_tp_size": self.tp_size, + "prefill_dp_size": self.dp_size, + } + return web.json_response(prefill_parallel_info, status=200) + + # Find corresponding prefill info + async with self.lock: + bootstrap_info = self.prefill_port_table[int(target_dp_group)][ + int(engine_rank) + ] + + if bootstrap_info is not None: + return web.json_response(bootstrap_info, status=200) + else: + return web.Response(text="Bootstrap info not Found", status=404) + + def _run_server(self): + try: + # Event Loop + self._loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._loop) + + self._runner = web.AppRunner(self.app) + self._loop.run_until_complete(self._runner.setup()) + + site = web.TCPSite(self._runner, host=self.host, port=self.port) + self._loop.run_until_complete(site.start()) + self._loop.run_forever() + except Exception as e: + logger.error(f"Server error: {str(e)}") + finally: + # Cleanup + self._loop.run_until_complete(self._runner.cleanup()) + self._loop.close() + + def close(self): + """Shutdown""" + if self._loop is not None and self._loop.is_running(): + self._loop.call_soon_threadsafe(self._loop.stop) + logger.info("Stopping server loop...") + + if self.thread.is_alive(): + self.thread.join(timeout=2) + logger.info("Server thread stopped") + + def poll(self) -> KVPoll: ... diff --git a/python/sglang/srt/disaggregation/common/utils.py b/python/sglang/srt/disaggregation/common/utils.py new file mode 100644 index 000000000..6f3da2128 --- /dev/null +++ b/python/sglang/srt/disaggregation/common/utils.py @@ -0,0 +1,42 @@ +import threading +from collections import deque +from typing import List, Tuple + +import numpy as np +import numpy.typing as npt + + +class FastQueue: + def __init__(self): + self._buf = deque() + self._cond = threading.Condition() + + def put(self, item): + with self._cond: + self._buf.append(item) + # wake up a thread of wait() + self._cond.notify() + + def get(self): + with self._cond: + # if queue is empty ,block until is notified() + while not self._buf: + self._cond.wait() + return self._buf.popleft() + + +def group_concurrent_contiguous( + src_indices: npt.NDArray[np.int32], dst_indices: npt.NDArray[np.int32] +) -> Tuple[List[npt.NDArray[np.int32]], List[npt.NDArray[np.int32]]]: + """Vectorised NumPy implementation.""" + if src_indices.size == 0: + return [], [] + + brk = np.where((np.diff(src_indices) != 1) | (np.diff(dst_indices) != 1))[0] + 1 + src_groups = np.split(src_indices, brk) + dst_groups = np.split(dst_indices, brk) + + src_groups = [g.tolist() for g in src_groups] + dst_groups = [g.tolist() for g in dst_groups] + + return src_groups, dst_groups diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py new file mode 100644 index 000000000..b79c8ca87 --- /dev/null +++ b/python/sglang/srt/disaggregation/decode.py @@ -0,0 +1,894 @@ +""" +Life cycle of a request in the decode server + +1. PreallocQueue: + a. Initialize a receiver for each request + b. The request handshakes first, and pre-allocate kv once there is available kv. + c. Move the request to TransferQueue. + +2. TransferQueue: + a. Poll the receiver to check the transfer state + b. If the transfer has finished, move the request to waiting queue + +3. WaitingQueue: + a. Use the requests in the queue to construct a PrebuiltExtendBatch + b. Skip the prefill forward but only populate metadata + +4. RunningBatch: + a. Merge the resolved PrebuiltExtendBatch into running batch to run decoding +""" + +from __future__ import annotations + +import logging +from collections import deque +from dataclasses import dataclass +from http import HTTPStatus +from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union + +import torch +from torch.distributed import ProcessGroup + +from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE +from sglang.srt.disaggregation.base import BaseKVManager, BaseKVReceiver, KVPoll +from sglang.srt.disaggregation.utils import ( + FAKE_BOOTSTRAP_HOST, + DisaggregationMode, + KVClassType, + MetadataBuffers, + ReqToMetadataIdxAllocator, + TransferBackend, + get_kv_class, + is_mla_backend, + kv_to_page_indices, + poll_and_all_reduce, + prepare_abort, +) +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.managers.schedule_batch import FINISH_ABORT, ScheduleBatch +from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache +from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter +from sglang.srt.utils import get_int_env_var, require_mlp_sync + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.managers.scheduler import Scheduler + +CLIP_MAX_NEW_TOKEN = get_int_env_var("SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION", 4096) + + +class DecodeReqToTokenPool: + """ + The difference of DecodeReqToTokenPool and ReqToTokenPool is that + DecodeReqToTokenPool subscribes memory for pre-allocated requests. + + In ReqToTokenPool, if `--max-running-requests` is 8, + #pre-allocated + #transfer + #running <= 8, but there are in fact more memory can carry pre-allocated requests. + + In DecodeReqToTokenPool, if `--max-running-requests` is 8, + #running <= 8, #pre-allocated + #transfer <= pre_alloc_size, so we can use the free memory to pre-allocate requests to unblock prefill. + """ + + def __init__( + self, + size: int, + max_context_len: int, + device: str, + enable_memory_saver: bool, + pre_alloc_size: int, + ): + memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=enable_memory_saver + ) + + self.size = size + self.max_context_len = max_context_len + self.device = device + self.pre_alloc_size = pre_alloc_size + with memory_saver_adapter.region(tag=GPU_MEMORY_TYPE_KV_CACHE): + self.req_to_token = torch.zeros( + (size + pre_alloc_size, max_context_len), + dtype=torch.int32, + device=device, + ) + + self.free_slots = list(range(size + pre_alloc_size)) + + def write(self, indices, values): + self.req_to_token[indices] = values + + def available_size(self): + return len(self.free_slots) + + def alloc(self, need_size: int) -> List[int]: + if need_size > len(self.free_slots): + return None + + select_index = self.free_slots[:need_size] + self.free_slots = self.free_slots[need_size:] + return select_index + + def free(self, free_index: Union[int, List[int]]): + if isinstance(free_index, (int,)): + self.free_slots.append(free_index) + else: + self.free_slots.extend(free_index) + + def clear(self): + self.free_slots = list(range(self.size + self.pre_alloc_size)) + + +@dataclass +class DecodeRequest: + req: Req + kv_receiver: BaseKVReceiver + waiting_for_input: bool = False + metadata_buffer_index: int = -1 + + +class DecodePreallocQueue: + """ + Store the requests that are preallocating. + """ + + def __init__( + self, + req_to_token_pool: ReqToTokenPool, + token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, + draft_token_to_kv_pool: Optional[KVCache], + req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator, + metadata_buffers: MetadataBuffers, + scheduler: Scheduler, + transfer_queue: DecodeTransferQueue, + tree_cache: BasePrefixCache, + gloo_group: ProcessGroup, + tp_rank: int, + tp_size: int, + dp_size: int, + gpu_id: int, + bootstrap_port: int, + max_total_num_tokens: int, + prefill_pp_size: int, + num_reserved_decode_tokens: int, + transfer_backend: TransferBackend, + ): + self.req_to_token_pool = req_to_token_pool + self.token_to_kv_pool_allocator = token_to_kv_pool_allocator + self.token_to_kv_pool = token_to_kv_pool_allocator.get_kvcache() + self.draft_token_to_kv_pool = draft_token_to_kv_pool + self.is_mla_backend = is_mla_backend(self.token_to_kv_pool) + self.metadata_buffers = metadata_buffers + self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator + self.scheduler = scheduler + self.transfer_queue = transfer_queue + self.tree_cache = tree_cache # this is always a chunk cache + self.gloo_group = gloo_group + self.tp_rank = tp_rank + self.tp_size = tp_size + self.dp_size = dp_size + self.gpu_id = gpu_id + self.bootstrap_port = bootstrap_port + self.max_total_num_tokens = max_total_num_tokens + self.prefill_pp_size = prefill_pp_size + self.num_reserved_decode_tokens = num_reserved_decode_tokens + self.transfer_backend = transfer_backend + # Queue for requests pending pre-allocation + self.queue: List[DecodeRequest] = [] + self.retracted_queue: List[Req] = [] + self.prefill_pp_size = prefill_pp_size + self.kv_manager = self._init_kv_manager() + + def _init_kv_manager(self) -> BaseKVManager: + kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS) + kv_args = kv_args_class() + + attn_tp_size = get_attention_tp_size() + kv_args.engine_rank = self.tp_rank % (attn_tp_size) + + kv_args.decode_tp_size = attn_tp_size + # Note(shangming): pp is not supported on the decode side yet, so its rank is fixed to 0 + kv_args.pp_rank = 0 + kv_args.system_dp_rank = self.scheduler.dp_rank + kv_args.prefill_pp_size = self.prefill_pp_size + kv_data_ptrs, kv_data_lens, kv_item_lens = ( + self.token_to_kv_pool.get_contiguous_buf_infos() + ) + if self.draft_token_to_kv_pool is not None: + # We should also transfer draft model kv cache. The indices are + # always shared with a target model. + draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( + self.draft_token_to_kv_pool.get_contiguous_buf_infos() + ) + kv_data_ptrs += draft_kv_data_ptrs + kv_data_lens += draft_kv_data_lens + kv_item_lens += draft_kv_item_lens + + kv_args.kv_data_ptrs = kv_data_ptrs + kv_args.kv_data_lens = kv_data_lens + kv_args.kv_item_lens = kv_item_lens + + kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = ( + self.metadata_buffers.get_buf_infos() + ) + + kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device + kv_args.gpu_id = self.scheduler.gpu_id + kv_manager_class: Type[BaseKVManager] = get_kv_class( + self.transfer_backend, KVClassType.MANAGER + ) + kv_manager: BaseKVManager = kv_manager_class( + kv_args, + DisaggregationMode.DECODE, + self.scheduler.server_args, + self.is_mla_backend, + ) + return kv_manager + + def add(self, req: Req, is_retracted: bool = False) -> None: + """Add a request to the pending queue.""" + if self._check_if_req_exceed_kv_capacity(req): + return + + if is_retracted: + self.retracted_queue.append(req) + else: + if req.bootstrap_host == FAKE_BOOTSTRAP_HOST: + kv_receiver_class = get_kv_class( + TransferBackend.FAKE, KVClassType.RECEIVER + ) + else: + kv_receiver_class = get_kv_class( + self.transfer_backend, KVClassType.RECEIVER + ) + + kv_receiver = kv_receiver_class( + mgr=self.kv_manager, + bootstrap_addr=f"{req.bootstrap_host}:{req.bootstrap_port}", + bootstrap_room=req.bootstrap_room, + prefill_dp_rank=req.data_parallel_rank, + ) + + self.queue.append( + DecodeRequest(req=req, kv_receiver=kv_receiver, waiting_for_input=False) + ) + + def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: + if len(req.origin_input_ids) > self.max_total_num_tokens: + message = f"Request {req.rid} exceeds the maximum number of tokens: {len(req.origin_input_ids)} > {self.max_total_num_tokens}" + logger.error(message) + prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) + self.scheduler.stream_output([req], req.return_logprob) + return True + return False + + def extend(self, reqs: List[Req], is_retracted: bool = False) -> None: + """Add a request to the pending queue.""" + for req in reqs: + self.add(req, is_retracted=is_retracted) + + def resume_retracted_reqs(self) -> List[Req]: + # TODO refactor the scheduling part, reuse with the unified engine logic as much as possible + + # allocate memory + resumed_reqs = [] + indices_to_remove = set() + allocatable_tokens = self._allocatable_tokens(count_retracted=False) + + for i, req in enumerate(self.retracted_queue): + if self.req_to_token_pool.available_size() <= 0: + break + + required_tokens_for_request = ( + len(req.origin_input_ids) + + len(req.output_ids) + + self.num_reserved_decode_tokens + ) + if required_tokens_for_request > allocatable_tokens: + break + + resumed_reqs.append(req) + indices_to_remove.add(i) + req.is_retracted = False + self._pre_alloc(req) + allocatable_tokens -= required_tokens_for_request + + # load from cpu, release the cpu copy + req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator) + + self.retracted_queue = [ + entry + for i, entry in enumerate(self.retracted_queue) + if i not in indices_to_remove + ] + + return resumed_reqs + + def _update_handshake_waiters(self) -> None: + if not self.queue: + return + + if all(decode_req.waiting_for_input for decode_req in self.queue): + return + + polls = poll_and_all_reduce( + [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group + ) + + for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): + if poll == KVPoll.Bootstrapping: + pass + elif poll == KVPoll.WaitingForInput: + decode_req.waiting_for_input = True + elif poll == KVPoll.Failed: + error_message = f"Decode handshake failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}" + try: + decode_req.kv_receiver.failure_exception() + except Exception as e: + error_message += f" with exception {e}" + logger.error(error_message) + prepare_abort( + decode_req.req, + error_message, + status_code=HTTPStatus.INTERNAL_SERVER_ERROR, + ) + if self.scheduler.enable_metrics: + self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() + else: + raise ValueError(f"Unexpected poll case: {poll}") + + def pop_preallocated(self) -> List[DecodeRequest]: + """Pop the preallocated requests from the pending queue (FIFO).""" + self._update_handshake_waiters() + + preallocated_reqs = [] + indices_to_remove = set() + + # We need to make sure that the sum of inflight tokens and allocatable tokens is greater than maximum input+output length of each inflight request + # Otherwise it is possible for one request running decode out of memory, while all other requests are in the transfer queue that cannot be retracted. + retractable_tokens = sum( + len(r.origin_input_ids) + len(r.output_ids) + for r in self.scheduler.running_batch.reqs + ) + allocatable_tokens = self._allocatable_tokens( + retractable_tokens=retractable_tokens, count_retracted=True + ) + # First, remove all failed requests from the queue + for i, decode_req in enumerate(self.queue): + if isinstance(decode_req.req.finished_reason, FINISH_ABORT): + self.scheduler.stream_output( + [decode_req.req], decode_req.req.return_logprob + ) + indices_to_remove.add(i) + + # Then, preallocate the remaining requests if possible + for i, decode_req in enumerate(self.queue): + if i in indices_to_remove: + continue + + if not decode_req.waiting_for_input: + continue + + if self.req_to_token_pool.available_size() <= 0: + break + + if self.req_to_metadata_buffer_idx_allocator.available_size() <= 0: + break + + # Memory estimation: don't add if the projected memory cannot be met + # TODO: add new_token ratio + origin_input_len = len(decode_req.req.origin_input_ids) + required_tokens_for_request = ( + origin_input_len + self.num_reserved_decode_tokens + ) + + if ( + max( + required_tokens_for_request, + origin_input_len + + min( + decode_req.req.sampling_params.max_new_tokens, + CLIP_MAX_NEW_TOKEN, + ) + - retractable_tokens, + ) + > allocatable_tokens + ): + break + if required_tokens_for_request > allocatable_tokens: + break + + allocatable_tokens -= required_tokens_for_request + self._pre_alloc(decode_req.req) + + kv_indices = ( + self.req_to_token_pool.req_to_token[decode_req.req.req_pool_idx][ + : len(decode_req.req.origin_input_ids) + ] + .cpu() + .numpy() + ) + + decode_req.metadata_buffer_index = ( + self.req_to_metadata_buffer_idx_allocator.alloc() + ) + assert decode_req.metadata_buffer_index is not None + page_indices = kv_to_page_indices( + kv_indices, self.token_to_kv_pool_allocator.page_size + ) + decode_req.kv_receiver.init(page_indices, decode_req.metadata_buffer_index) + preallocated_reqs.append(decode_req) + indices_to_remove.add(i) + + self.queue = [ + entry for i, entry in enumerate(self.queue) if i not in indices_to_remove + ] + + return preallocated_reqs + + @property + def num_tokens_pre_allocated(self): + return sum( + len(decode_req.req.fill_ids) for decode_req in self.transfer_queue.queue + ) + + def _allocatable_tokens( + self, retractable_tokens: Optional[int] = None, count_retracted: bool = True + ) -> int: + need_space_for_single_req = ( + max( + [ + min(x.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKEN) + + len(x.origin_input_ids) + - retractable_tokens + for x in self.scheduler.running_batch.reqs + ] + ) + if retractable_tokens is not None + and len(self.scheduler.running_batch.reqs) > 0 + else 0 + ) + + if self.scheduler.model_config.is_hybrid: + available_size = min( + self.token_to_kv_pool_allocator.full_available_size(), + self.token_to_kv_pool_allocator.swa_available_size(), + ) + else: + available_size = self.token_to_kv_pool_allocator.available_size() + + allocatable_tokens = available_size - max( + # preserve some space for future decode + self.num_reserved_decode_tokens + * ( + len(self.scheduler.running_batch.reqs) + + len(self.transfer_queue.queue) + + len(self.scheduler.waiting_queue) + ), + # make sure each request can finish if reach max_tokens with all other requests retracted + need_space_for_single_req, + ) + + # Note: if the last fake extend just finishes, and we enter `pop_preallocated` immediately in the next iteration + # the extend batch is not in any queue, so we need to explicitly add the tokens slots here + if ( + self.scheduler.last_batch + and self.scheduler.last_batch.forward_mode.is_extend() + ): + allocatable_tokens -= self.num_reserved_decode_tokens * len( + self.scheduler.last_batch.reqs + ) + + if count_retracted: + allocatable_tokens -= sum( + [ + len(req.origin_input_ids) + + len(req.output_ids) + + self.num_reserved_decode_tokens + for req in self.retracted_queue + ] + ) + return allocatable_tokens + + def _pre_alloc(self, req: Req) -> torch.Tensor: + """Pre-allocate the memory for req_to_token and token_kv_pool""" + req_pool_indices = self.req_to_token_pool.alloc(1) + + assert ( + req_pool_indices is not None + ), "req_pool_indices is full! There is a bug in memory estimation." + + req.req_pool_idx = req_pool_indices[0] + + if self.token_to_kv_pool_allocator.page_size == 1: + kv_loc = self.token_to_kv_pool_allocator.alloc( + len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0) + ) + else: + num_tokens = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0) + kv_loc = self.token_to_kv_pool_allocator.alloc_extend( + prefix_lens=torch.tensor( + [0], + dtype=torch.int64, + device=self.token_to_kv_pool_allocator.device, + ), + seq_lens=torch.tensor( + [num_tokens], + dtype=torch.int64, + device=self.token_to_kv_pool_allocator.device, + ), + last_loc=torch.tensor( + [-1], + dtype=torch.int64, + device=self.token_to_kv_pool_allocator.device, + ), + extend_num_tokens=num_tokens, + ) + + assert ( + kv_loc is not None + ), "KV cache is full! There is a bug in memory estimation." + + self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc) + + # populate metadata + req.fill_ids = req.origin_input_ids + req.output_ids + req.extend_input_len = len(req.origin_input_ids) + + return kv_loc + + +class DecodeTransferQueue: + """ + Store the requests that is polling kv + """ + + def __init__( + self, + gloo_group: ProcessGroup, + req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator, + tp_rank: int, + metadata_buffers: MetadataBuffers, + scheduler: Scheduler, + tree_cache: BasePrefixCache, + ): + self.queue: List[DecodeRequest] = [] + self.gloo_group = gloo_group + self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator + self.tp_rank = tp_rank + self.metadata_buffers = metadata_buffers + self.scheduler = scheduler + self.tree_cache = tree_cache + self.spec_algorithm = scheduler.spec_algorithm + + def add(self, decode_req: DecodeRequest) -> None: + self.queue.append(decode_req) + + def extend(self, decode_reqs: List[DecodeRequest]) -> None: + self.queue.extend(decode_reqs) + + def pop_transferred(self) -> List[Req]: + if not self.queue: + return [] + polls = poll_and_all_reduce( + [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group + ) + + transferred_reqs = [] + indices_to_remove = set() + for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): + if poll == KVPoll.Failed: + error_message = f"Decode transfer failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}" + try: + decode_req.kv_receiver.failure_exception() + except Exception as e: + error_message += f" with exception {e}" + logger.error(error_message) + prepare_abort( + decode_req.req, + error_message, + status_code=HTTPStatus.INTERNAL_SERVER_ERROR, + ) + self.scheduler.stream_output( + [decode_req.req], decode_req.req.return_logprob + ) + # unlock the kv cache or it will have memory leak + self.tree_cache.cache_finished_req(decode_req.req) + indices_to_remove.add(i) + if self.scheduler.enable_metrics: + self.scheduler.metrics_collector.increment_transfer_failed_reqs() + continue + elif poll == KVPoll.Success: + + idx = decode_req.metadata_buffer_index + ( + output_id, + output_token_logprobs_val, + output_token_logprobs_idx, + output_top_logprobs_val, + output_top_logprobs_idx, + output_hidden_states, + ) = self.metadata_buffers.get_buf(idx) + + decode_req.req.output_ids.append(output_id[0].item()) + if not self.spec_algorithm.is_none(): + decode_req.req.hidden_states_tensor = output_hidden_states + if decode_req.req.return_logprob: + decode_req.req.output_token_logprobs_val.append( + output_token_logprobs_val[0].item() + ) + decode_req.req.output_token_logprobs_idx.append( + output_token_logprobs_idx[0].item() + ) + decode_req.req.output_top_logprobs_val.append( + output_top_logprobs_val[ + : decode_req.req.top_logprobs_num + ].tolist() + ) + decode_req.req.output_top_logprobs_idx.append( + output_top_logprobs_idx[ + : decode_req.req.top_logprobs_num + ].tolist() + ) + + if hasattr(decode_req.kv_receiver, "clear"): + decode_req.kv_receiver.clear() + + # special handling for sampling_params.max_new_tokens == 1 + if decode_req.req.sampling_params.max_new_tokens == 1: + # finish immediately + decode_req.req.check_finished() + self.scheduler.stream_output( + [decode_req.req], decode_req.req.return_logprob + ) + self.tree_cache.cache_finished_req(decode_req.req) + else: + transferred_reqs.append(decode_req.req) + + indices_to_remove.add(i) + elif poll in [ + KVPoll.Bootstrapping, + KVPoll.WaitingForInput, + KVPoll.Transferring, + ]: + pass + else: + raise ValueError(f"Unexpected poll case: {poll}") + + for i in indices_to_remove: + idx = self.queue[i].metadata_buffer_index + assert idx != -1 + self.req_to_metadata_buffer_idx_allocator.free(idx) + + self.queue = [ + entry for i, entry in enumerate(self.queue) if i not in indices_to_remove + ] + + return transferred_reqs + + +class SchedulerDisaggregationDecodeMixin: + + @torch.no_grad() + def event_loop_normal_disagg_decode(self: Scheduler): + """A normal scheduler loop for decode worker in disaggregation mode.""" + + while True: + recv_reqs = self.recv_requests() + self.process_input_requests(recv_reqs) + # polling and allocating kv cache + self.process_decode_queue() + batch = self.get_next_disagg_decode_batch_to_run() + self.cur_batch = batch + + prepare_mlp_sync_flag = require_mlp_sync(self.server_args) + + if batch: + # Generate fake extend output. + if batch.forward_mode.is_extend(): + # Note: Logprobs should be handled on the prefill engine. + self.stream_output( + batch.reqs, any(req.return_logprob for req in batch.reqs) + ) + if prepare_mlp_sync_flag: + self._prepare_idle_batch_and_run(None) + else: + if prepare_mlp_sync_flag: + self.prepare_mlp_sync_batch(batch) + result = self.run_batch(batch) + self.process_batch_result(batch, result) + elif prepare_mlp_sync_flag: + batch, _ = self._prepare_idle_batch_and_run(None) + + if batch is None and ( + len(self.waiting_queue) + + len(self.disagg_decode_transfer_queue.queue) + + len(self.disagg_decode_prealloc_queue.queue) + == 0 + ): + self.self_check_during_idle() + + self.last_batch = batch + + @torch.no_grad() + def event_loop_overlap_disagg_decode(self: Scheduler): + result_queue = deque() + self.last_batch: Optional[ScheduleBatch] = None + self.last_batch_in_queue = False # last batch is modified in-place, so we need another variable to track if it's extend + + while True: + recv_reqs = self.recv_requests() + self.process_input_requests(recv_reqs) + # polling and allocating kv cache + self.process_decode_queue() + batch = self.get_next_disagg_decode_batch_to_run() + self.cur_batch = batch + last_batch_in_queue = False + + prepare_mlp_sync_flag = require_mlp_sync(self.server_args) + + if batch: + # Generate fake extend output. + if batch.forward_mode.is_extend(): + # Note: Logprobs should be handled on the prefill engine. + self.stream_output( + batch.reqs, any(req.return_logprob for req in batch.reqs) + ) + if prepare_mlp_sync_flag: + batch_, result = self._prepare_idle_batch_and_run( + None, delay_process=True + ) + if batch_: + result_queue.append((batch_.copy(), result)) + last_batch_in_queue = True + else: + if prepare_mlp_sync_flag: + self.prepare_mlp_sync_batch(batch) + result = self.run_batch(batch) + result_queue.append((batch.copy(), result)) + + if (self.last_batch is None) or (not self.last_batch_in_queue): + # Create a dummy first batch to start the pipeline for overlap schedule. + # It is now used for triggering the sampling_info_done event. + tmp_batch = ScheduleBatch( + reqs=None, + forward_mode=ForwardMode.DUMMY_FIRST, + next_batch_sampling_info=self.tp_worker.cur_sampling_info, + ) + self.set_next_batch_sampling_info_done(tmp_batch) + last_batch_in_queue = True + + elif prepare_mlp_sync_flag: + batch, result = self._prepare_idle_batch_and_run( + None, delay_process=True + ) + if batch: + result_queue.append((batch.copy(), result)) + last_batch_in_queue = True + + # Process the results of the previous batch but skip if the last batch is extend + if self.last_batch and self.last_batch_in_queue: + tmp_batch, tmp_result = result_queue.popleft() + tmp_batch.next_batch_sampling_info = ( + self.tp_worker.cur_sampling_info if batch else None + ) + self.process_batch_result(tmp_batch, tmp_result) + + if batch is None and ( + len(self.waiting_queue) + + len(self.disagg_decode_transfer_queue.queue) + + len(self.disagg_decode_prealloc_queue.queue) + == 0 + ): + self.self_check_during_idle() + + self.last_batch = batch + self.last_batch_in_queue = last_batch_in_queue + + def _prepare_idle_batch_and_run(self: Scheduler, batch, delay_process=False): + batch = self.prepare_mlp_sync_batch(batch) + result = None + if batch: + result = self.run_batch(batch) + if not delay_process: + self.process_batch_result(batch, result) + return batch, result + + def get_next_disagg_decode_batch_to_run( + self: Scheduler, + ) -> Optional[Tuple[ScheduleBatch, bool]]: + """Create fake completed prefill if possible and merge with running batch""" + # Merge the prefill batch into the running batch + last_batch = self.last_batch + if last_batch and last_batch.forward_mode.is_extend(): + # chunked prefill doesn't happen in decode instance. + assert self.chunked_req is None + # Filter finished batches. + last_batch.filter_batch() + if not last_batch.is_empty(): + if self.running_batch.is_empty(): + self.running_batch = last_batch + else: + # merge running_batch with prefill batch + self.running_batch.merge_batch(last_batch) + + new_prebuilt_batch = self.get_new_prebuilt_batch() + + ret: Optional[ScheduleBatch] = None + if new_prebuilt_batch: + ret = new_prebuilt_batch + else: + if self.running_batch.is_empty(): + ret = None + else: + self.running_batch = self.update_running_batch(self.running_batch) + ret = self.running_batch if not self.running_batch.is_empty() else None + + return ret + + def get_new_prebuilt_batch(self: Scheduler) -> Optional[ScheduleBatch]: + """Create a schedulebatch for fake completed prefill""" + if self.grammar_queue: + self.move_ready_grammar_requests() + + if len(self.waiting_queue) == 0: + return None + + curr_batch_size = self.running_batch.batch_size() + + batch_size = min(self.req_to_token_pool.size, self.max_running_requests) + + num_not_used_batch = batch_size - curr_batch_size + + # pop req from waiting queue + can_run_list: List[Req] = [] + waiting_queue: List[Req] = [] + + for i in range(len(self.waiting_queue)): + req = self.waiting_queue[i] + # we can only add at least `num_not_used_batch` new batch to the running queue + if i < num_not_used_batch: + can_run_list.append(req) + req.init_next_round_input(self.tree_cache) + else: + waiting_queue.append(req) + + self.waiting_queue = waiting_queue + if len(can_run_list) == 0: + return None + + # construct a schedule batch with those requests and mark as decode + new_batch = ScheduleBatch.init_new( + can_run_list, + self.req_to_token_pool, + self.token_to_kv_pool_allocator, + self.tree_cache, + self.model_config, + self.enable_overlap, + self.spec_algorithm, + ) + + # construct fake completed prefill + new_batch.prepare_for_prebuilt_extend() + new_batch.process_prebuilt_extend(self.server_args, self.model_config) + + return new_batch + + def process_decode_queue(self: Scheduler): + # try to resume retracted requests if there are enough space for another `num_reserved_decode_tokens` decode steps + resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs() + self.waiting_queue.extend(resumed_reqs) + if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0: + # if there are still retracted requests, we do not allocate new requests + return + + req_conns = self.disagg_decode_prealloc_queue.pop_preallocated() + self.disagg_decode_transfer_queue.extend(req_conns) + alloc_reqs = ( + self.disagg_decode_transfer_queue.pop_transferred() + ) # the requests which kv has arrived + self.waiting_queue.extend(alloc_reqs) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py new file mode 100644 index 000000000..c1cb17c04 --- /dev/null +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +import logging +from http import HTTPStatus +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.disaggregation.utils import prepare_abort +from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode +from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.server_args import ServerArgs + + +class ScheduleBatchDisaggregationDecodeMixin: + + def prepare_for_prebuilt_extend(self: ScheduleBatch): + """ + Prepare a prebuilt extend by populate metadata + Adapted from .prepare_for_extend(). + """ + + self.forward_mode = ForwardMode.EXTEND + reqs = self.reqs + input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs] + extend_num_tokens = sum(len(ids) for ids in input_ids) + seq_lens = [] + pre_lens = [] + req_pool_indices = [] + + # Pre-calculate total size + total_size = sum(req.extend_input_len for req in reqs) + out_cache_loc = torch.empty(total_size, dtype=torch.int64, device=self.device) + + # Fill the tensor in one pass + offset = 0 + for i, req in enumerate(reqs): + req_pool_indices.append(req.req_pool_idx) + + chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][ + : req.extend_input_len + ] + assert ( + offset + req.extend_input_len <= total_size + ), f"Exceeds total size: offset={offset}, req.extend_input_len={req.extend_input_len}, total_size={total_size}" + out_cache_loc[offset : offset + req.extend_input_len] = chunk + offset += req.extend_input_len + + pre_len = len(req.prefix_indices) + seq_len = len(req.origin_input_ids) + max(0, len(req.output_ids) - 1) + seq_lens.append(seq_len) + if len(req.output_ids) == 0: + assert ( + seq_len - pre_len == req.extend_input_len + ), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_input_len={req.extend_input_len}" + + req.cached_tokens += pre_len - req.already_computed + req.already_computed = seq_len + req.is_retracted = False + pre_lens.append(pre_len) + req.extend_logprob_start_len = 0 + + extend_input_logprob_token_ids = None + + # Set fields + self.input_ids = torch.tensor( + sum(input_ids, []), dtype=torch.int32, device=self.device + ) + self.req_pool_indices = torch.tensor( + req_pool_indices, dtype=torch.int64, device=self.device + ) + self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, device=self.device) + self.orig_seq_lens = torch.tensor( + seq_lens, dtype=torch.int32, device=self.device + ) + self.out_cache_loc = out_cache_loc + self.seq_lens_sum = sum(seq_lens) + + if self.return_logprob: + self.top_logprobs_nums = [r.top_logprobs_num for r in reqs] + self.token_ids_logprobs = [r.token_ids_logprob for r in reqs] + + self.extend_num_tokens = extend_num_tokens + self.prefix_lens = [len(r.prefix_indices) for r in reqs] + self.extend_lens = [r.extend_input_len for r in reqs] + self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs] + self.extend_input_logprob_token_ids = extend_input_logprob_token_ids + self.multimodal_inputs = [r.multimodal_inputs for r in reqs] + + # Build sampling info + self.sampling_info = SamplingBatchInfo.from_schedule_batch( + self, + self.model_config.vocab_size, + ) + + def process_prebuilt_extend( + self: ScheduleBatch, server_args: ServerArgs, model_config: ModelConfig + ): + """Assign the buffered last input id to schedule batch""" + self.output_ids = [] + for req in self.reqs: + self.output_ids.append(req.output_ids[-1]) + self.tree_cache.cache_unfinished_req(req) + if req.grammar is not None: + # FIXME: this try-except block is for handling unexpected xgrammar issue. + try: + req.grammar.accept_token(req.output_ids[-1]) + except ValueError as e: + # Grammar accept_token can raise ValueError if the token is not in the grammar. + # This can happen if the grammar is not set correctly or the token is invalid. + error_message = f"Grammar accept_token failed for req {req.rid} with token {req.output_ids[-1]}: {e}" + self.tree_cache.cache_finished_req(req) + prepare_abort( + req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR + ) + req.grammar.finished = req.finished() + self.output_ids = torch.tensor(self.output_ids, device=self.device) + + # Simulate the eagle run. We add mock data to hidden states for the + # ease of implementation now meaning the first token will have acc rate + # of 0. + if not self.spec_algorithm.is_none(): + + b = len(self.reqs) + topk_p = torch.arange( + b * server_args.speculative_eagle_topk, + 0, + -1, + device=self.device, + dtype=torch.float32, + ) + topk_p = topk_p.reshape(b, server_args.speculative_eagle_topk) + topk_p /= b * server_args.speculative_eagle_topk + topk_index = torch.arange( + b * server_args.speculative_eagle_topk, device=self.device + ) + topk_index = topk_index.reshape(b, server_args.speculative_eagle_topk) + + hidden_states_list = [req.hidden_states_tensor for req in self.reqs] + hidden_states = torch.stack(hidden_states_list, dim=0).to(self.device) + + # local import to avoid circular import + from sglang.srt.speculative.eagle_utils import EagleDraftInput + + spec_info = EagleDraftInput( + topk_p=topk_p, + topk_index=topk_index, + hidden_states=hidden_states, + verified_id=self.output_ids, + ) + spec_info.prepare_for_extend(self) + spec_info.capture_hidden_mode = CaptureHiddenMode.LAST + self.spec_info = spec_info diff --git a/python/sglang/srt/disaggregation/fake/__init__.py b/python/sglang/srt/disaggregation/fake/__init__.py new file mode 100644 index 000000000..d7cdb4b27 --- /dev/null +++ b/python/sglang/srt/disaggregation/fake/__init__.py @@ -0,0 +1 @@ +from sglang.srt.disaggregation.fake.conn import FakeKVReceiver, FakeKVSender diff --git a/python/sglang/srt/disaggregation/fake/conn.py b/python/sglang/srt/disaggregation/fake/conn.py new file mode 100644 index 000000000..120633824 --- /dev/null +++ b/python/sglang/srt/disaggregation/fake/conn.py @@ -0,0 +1,85 @@ +import logging +from typing import List, Optional + +import numpy as np +import numpy.typing as npt + +from sglang.srt.disaggregation.base.conn import ( + BaseKVManager, + BaseKVReceiver, + BaseKVSender, + KVPoll, +) + +logger = logging.getLogger(__name__) + + +# For warmup reqs, we don't kv transfer, we use the fake sender and receiver +class FakeKVSender(BaseKVSender): + def __init__( + self, + mgr: BaseKVManager, + bootstrap_addr: str, + bootstrap_room: int, + dest_tp_ranks: List[int], + pp_rank: int, + ): + self.has_sent = False + + def poll(self) -> KVPoll: + if self.has_sent is False: + # Assume handshake completed instantly + return KVPoll.WaitingForInput + else: + # Assume transfer completed instantly + logger.debug("FakeKVSender poll success") + return KVPoll.Success + + def init( + self, + kv_indices: list[int], + aux_index: Optional[int] = None, + ): + logger.debug( + f"FakeKVSender init with kv_indices: {kv_indices}, aux_index: {aux_index}" + ) + pass + + def send( + self, + kv_indices: npt.NDArray[np.int32], + ): + self.has_sent = True + logger.debug(f"FakeKVSender send with kv_indices: {kv_indices}") + + def failure_exception(self): + raise Exception("Fake KVSender Exception") + + +class FakeKVReceiver(BaseKVReceiver): + def __init__( + self, + mgr: BaseKVManager, + bootstrap_addr: str, + bootstrap_room: Optional[int] = None, + prefill_dp_rank: Optional[int] = None, + ): + self.has_init = False + + def poll(self) -> KVPoll: + if self.has_init is False: + # Assume handshake completed instantly + return KVPoll.WaitingForInput + else: + # Assume transfer completed instantly + logger.debug("FakeKVReceiver poll success") + return KVPoll.Success + + def init(self, kv_indices: list[int], aux_index: Optional[int] = None): + self.has_init = True + logger.debug( + f"FakeKVReceiver init with kv_indices: {kv_indices}, aux_index: {aux_index}" + ) + + def failure_exception(self): + raise Exception("Fake KVReceiver Exception") diff --git a/python/sglang/srt/disaggregation/kv_events.py b/python/sglang/srt/disaggregation/kv_events.py new file mode 100644 index 000000000..f0a3a4357 --- /dev/null +++ b/python/sglang/srt/disaggregation/kv_events.py @@ -0,0 +1,412 @@ +""" +Copyright 2025 SGLang 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. +""" + +""" +KV caching events +""" + +import atexit +import logging +import queue +import threading +import time +from abc import ABC, abstractmethod +from collections import deque +from itertools import count +from queue import Queue +from typing import Any, Callable, Optional, Union + +import msgspec +import zmq +from pydantic import BaseModel + +logger = logging.getLogger(__name__) + + +class EventBatch( + msgspec.Struct, + array_like=True, # type: ignore[call-arg] + omit_defaults=True, # type: ignore[call-arg] + gc=False, # type: ignore[call-arg] +): + ts: float + events: list[Any] + attn_dp_rank: Optional[int] = None + + +class KVCacheEvent( + msgspec.Struct, + array_like=True, # type: ignore[call-arg] + omit_defaults=True, # type: ignore[call-arg] + gc=False, # type: ignore[call-arg] + tag=True, +): + """Base class for all KV cache-related events""" + + +class BlockStored(KVCacheEvent): + block_hashes: list[int] + parent_block_hash: Optional[int] + token_ids: list[int] + block_size: int + lora_id: Optional[int] + + +class BlockRemoved(KVCacheEvent): + block_hashes: list[int] + + +class AllBlocksCleared(KVCacheEvent): + pass + + +class KVEventBatch(EventBatch): + events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]] + + +class EventPublisher(ABC): + """ + Lightweight publisher for EventBatch batches with + support for DP attention. + + In DP attention - each rank has its own Scheduler and + KV cache instance in order to avoid duplicate events + and ensure proper event attribution. In our implementation + + - Each DP rank has its own EventPublisher + - Publishers annotate events with the dp rank + - This allows consumers to distinguish events from different DP ranks + """ + + def __init__(self, attn_dp_rank: int = 0): + self._attn_dp_rank = attn_dp_rank + + @abstractmethod + def publish(self, events: EventBatch) -> None: + """Emit events in order. + + Implementations should guarantee at-least-once delivery and + monotonic ordering (e.g., via sequence numbers). + """ + + @abstractmethod + def shutdown(self) -> None: + """Shutdown the publisher.""" + + +class NullEventPublisher(EventPublisher): + """No-op implementation (default when disabled).""" + + def publish(self, events) -> None: + return + + def shutdown(self) -> None: + return + + +class ZmqEventPublisher(EventPublisher): + """Reliable PUB/ROUTER publisher with an in-memory replay buffer. + + Spawns a separate thread to handle publishing from a queue. + + Parameters + ---------- + endpoint: + PUB address. Use ``tcp://*:5557`` to bind or ``tcp://host:5557`` to + connect. + replay_endpoint: + Optional ROUTER address for replay requests. When given, subscribers can + request missed batches by sending the starting sequence number as an + 8-byte big-endian integer. + buffer_steps: + Number of past batches to keep for replay. + hwm: + ZeroMQ high-water-mark for PUB socket. + max_queue_size: + Maximum number of events to buffer in memory. + topic: + Topic to publish events to. + """ + + SHUTDOWN_TIMEOUT: float = 1.0 + END_SEQ = (-1).to_bytes(8, "big", signed=True) + + def __init__( + self, + attn_dp_rank: int, + endpoint: str = "tcp://*:5557", + replay_endpoint: Optional[str] = None, + buffer_steps: int = 10_000, + hwm: int = 100_000, + max_queue_size: int = 100_000, + topic: str = "", + ) -> None: + # Storage + super().__init__(attn_dp_rank) + self._event_queue = Queue[Optional[EventBatch]](maxsize=max_queue_size) + self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps) + + # ZMQ sockets + self._ctx = zmq.Context.instance() + self._pub: Optional[zmq.Socket] = None + self._replay: Optional[zmq.Socket] = None + self._dp_rank = attn_dp_rank + self._endpoint = self.offset_endpoint_port(endpoint, self._dp_rank) + self._replay_endpoint = self.offset_endpoint_port( + replay_endpoint, self._dp_rank + ) + self._hwm = hwm + self._socket_setup() + + # Payload + self._seq_gen = count() + self._topic_bytes = topic.encode("utf-8") + + # Thread + self._running = True + logger.info("Starting ZMQ publisher thread") + + self._thread = threading.Thread( + target=self._publisher_thread, daemon=True, name="zmq-publisher" + ) + self._thread.start() + + atexit.register(self.shutdown) + + def publish(self, events: EventBatch) -> None: + if not self._running: + raise RuntimeError("Publisher is closed") + if events.attn_dp_rank is None: + events.attn_dp_rank = self._dp_rank + self._event_queue.put(events) + + def shutdown(self) -> None: + """Stop the publisher thread and clean up resources.""" + self._running = False + self._event_queue.put_nowait(None) + + start = time.time() + pending_items = True + while pending_items and (time.time() - start < self.SHUTDOWN_TIMEOUT): + pending_items = not self._event_queue.empty() + if pending_items: + time.sleep(0.1) + + if pending_items: + logger.warning( + "Warning: Queue still has %s items after %s seconds timeout", + self._event_queue.qsize(), + self.SHUTDOWN_TIMEOUT, + ) + + if self._thread.is_alive(): + self._thread.join(timeout=self.SHUTDOWN_TIMEOUT) + + # Clean up ZMQ resources + try: + if self._pub is not None: + self._pub.close(linger=0) + if self._replay is not None: + self._replay.close(linger=0) + finally: + pass # Do not terminate context; other sockets may use it + + def _socket_setup(self) -> None: + """Initialize sockets + https://pyzmq.readthedocs.io/en/v19.0.0/morethanbindings.html#thread-safety + """ + if self._pub is None: + self._pub = self._ctx.socket(zmq.PUB) + self._pub.set_hwm(self._hwm) + # Heuristic: bind if wildcard / * present, else connect. + # bind stable, connect volatile convention + if ( + "*" in self._endpoint + or "::" in self._endpoint + or self._endpoint.startswith("ipc://") + or self._endpoint.startswith("inproc://") + ): + self._pub.bind(self._endpoint) + else: + self._pub.connect(self._endpoint) + + # Set up replay socket: use ROUTER + # 1) handles multiple REQ clients (identities) + # 2) lets us send back one request → many replies (streamed events) + # 3) works in our non‑blocking poll loop alongside PUB + if self._replay_endpoint is not None: + self._replay = self._ctx.socket(zmq.ROUTER) + self._replay.bind(self._replay_endpoint) + + def _publisher_thread(self) -> None: + """Background thread that processes the event queue.""" + self._pack = msgspec.msgpack.Encoder() + + assert self._pub is not None # narrows type for mypy + + while self._running or self._event_queue.qsize() > 0: + # --- replay (non-critical) --------------------------------- + if self._replay is not None and self._replay.poll(0): + try: + self._service_replay() + except Exception as e: + logger.exception("Error in replay: %s", e) + + # --- main queue (critical) --------------------------------- + try: + event = self._event_queue.get(timeout=0.1) + if event is None: + break # Sentinel received, exit thread + except queue.Empty: + continue + + try: + seq = next(self._seq_gen) + + payload = self._pack.encode(event) + seq_bytes = seq.to_bytes(8, "big") + self._pub.send_multipart((self._topic_bytes, seq_bytes, payload)) + + self._buffer.append((seq, payload)) + self._event_queue.task_done() + + except Exception as e: + # Publishing failed; back-off a bit to avoid a tight error loop + logger.exception("Error in publisher thread: %s", e) + time.sleep(0.1) + + def _service_replay(self) -> None: + """If a replay request is waiting, send buffered batches.""" + assert self._replay is not None # narrows type for mypy + + frame = self._replay.recv_multipart() + if len(frame) != 3: + logger.warning("Invalid replay request: %s", frame) + return + client_id, _, start_seq_bytes = frame + start_seq = int.from_bytes(start_seq_bytes, "big") + + for seq, buf in self._buffer: + if seq >= start_seq: + # [identity, empty_delim, seq_bytes, payload] + # (identity, empty_delim) are stripped off by the router + # receiving payload is (seq_bytes, payload) + self._replay.send_multipart( + (client_id, b"", seq.to_bytes(8, "big"), buf) + ) + # Send end of sequence marker + # receiving payload is (-1, b""") + self._replay.send_multipart((client_id, b"", self.END_SEQ, b"")) + + @staticmethod + def offset_endpoint_port( + endpoint: Optional[str], data_parallel_rank: int + ) -> Optional[str]: + """Helper function to offset the port in an endpoint by + the data parallel rank. + + Args: + endpoint: The endpoint string + (e.g., "tcp://*:5557" or "inproc://cache") + data_parallel_rank: The data parallel rank to offset by + + Returns: + The endpoint with the port offset by data_parallel_rank + or suffix appended + """ + # Do nothing if input is None or data_parallel_rank is 0 + if not endpoint or data_parallel_rank == 0: + return endpoint + + if "inproc" in endpoint: + return f"{endpoint}_dp{data_parallel_rank}" + if "tcp" in endpoint: + if endpoint and ":" in endpoint: + # Get everything after the last colon (the port) + last_colon_idx = endpoint.rfind(":") + base_addr = endpoint[:last_colon_idx] + base_port = int(endpoint[last_colon_idx + 1 :]) + new_port = base_port + data_parallel_rank + return f"{base_addr}:{new_port}" + return endpoint + raise ValueError("Invalid endpoint: must contain 'inproc' or 'tcp'") + + +class KVEventsConfig(BaseModel): + """Configuration for KV event publishing.""" + + publisher: str = "null" + """The publisher to use for publishing kv events. Can be "null", "zmq". + """ + + endpoint: str = "tcp://*:5557" + """The zmq endpoint to use for publishing kv events. + """ + + replay_endpoint: Optional[str] = None + """The zmq endpoint to use for replaying kv events. + """ + + buffer_steps: int = 10_000 + """The number of steps to cache for replay endpoint. Will only save + events from the last N steps for the replay endpoint. + """ + + hwm: int = 100_000 + """The zmq high water mark for the event publisher. After queueing N events, + events will start dropping if the consumer is not keeping up. + """ + + max_queue_size: int = 100_000 + """The maximum number of events to queue while waiting for publishing. + """ + + topic: str = "" + """The topic to use for the event publisher. Consumers can subscribe to + this topic to receive events. + """ + + @classmethod + def from_cli(cls, cli_value: str) -> "KVEventsConfig": + """Parse the CLI value for the event publisher config.""" + return KVEventsConfig.model_validate_json(cli_value) + + +class EventPublisherFactory: + _registry: dict[str, Callable[..., EventPublisher]] = { + "null": NullEventPublisher, + "zmq": ZmqEventPublisher, + } + + @classmethod + def register_publisher(cls, name: str, ctor: Callable[..., EventPublisher]) -> None: + if name in cls._registry: + raise KeyError(f"publisher '{name}' already registered") + cls._registry[name] = ctor + + @classmethod + def create(cls, config: Optional[str], attn_dp_rank: int = 0) -> EventPublisher: + """Create publisher from a config mapping.""" + if not config: + return NullEventPublisher() + config = KVEventsConfig.from_cli(config) + config_dict = config.model_dump() + + kind = config_dict.pop("publisher", "null") + try: + constructor = cls._registry[kind] + except KeyError as exc: + raise ValueError(f"Unknown event publisher '{kind}'") from exc + return constructor(attn_dp_rank=attn_dp_rank, **config_dict) diff --git a/python/sglang/srt/disaggregation/mini_lb.py b/python/sglang/srt/disaggregation/mini_lb.py new file mode 100644 index 000000000..5aaa2a70e --- /dev/null +++ b/python/sglang/srt/disaggregation/mini_lb.py @@ -0,0 +1,6 @@ +raise RuntimeError( + """The 'mini_lb' module has been relocated to the 'sglang_router' package. + We recommend installing 'sglang-router' with Rust support for optimal performance. + If you encounter issues building the router with Rust, set the environment variable + 'SGLANG_ROUTER_BUILD_NO_RUST=1' and add '--mini-lb' to the command line to use the Python version of 'mini_lb'.""" +) diff --git a/python/sglang/srt/disaggregation/mooncake/__init__.py b/python/sglang/srt/disaggregation/mooncake/__init__.py new file mode 100644 index 000000000..bea967e4e --- /dev/null +++ b/python/sglang/srt/disaggregation/mooncake/__init__.py @@ -0,0 +1,6 @@ +from sglang.srt.disaggregation.mooncake.conn import ( + MooncakeKVBootstrapServer, + MooncakeKVManager, + MooncakeKVReceiver, + MooncakeKVSender, +) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py new file mode 100644 index 000000000..f69d29622 --- /dev/null +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -0,0 +1,1704 @@ +from __future__ import annotations + +import asyncio +import concurrent.futures +import ctypes +import dataclasses +import logging +import os +import queue +import socket +import struct +import threading +import time +from collections import defaultdict +from functools import cache +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import numpy.typing as npt +import requests +import zmq +from aiohttp import web + +from sglang.srt.disaggregation.base.conn import ( + BaseKVBootstrapServer, + BaseKVManager, + BaseKVReceiver, + BaseKVSender, + KVArgs, + KVPoll, +) +from sglang.srt.disaggregation.common.utils import ( + FastQueue, + group_concurrent_contiguous, +) +from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.distributed import get_pp_group +from sglang.srt.layers.dp_attention import ( + get_attention_dp_rank, + get_attention_dp_size, + get_attention_tp_rank, + get_attention_tp_size, +) +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import ( + format_tcp_address, + get_bool_env_var, + get_free_port, + get_int_env_var, + get_ip, + get_local_ip_auto, + is_valid_ipv6_address, + maybe_wrap_ipv6_address, +) + +logger = logging.getLogger(__name__) + + +class KVTransferError(Exception): + def __init__(self, bootstrap_room: int, failure_reason: str): + super().__init__(failure_reason) + self.bootstrap_room = bootstrap_room + self.failure_reason = failure_reason + + def __str__(self): + return f"KVTransferError(bootstrap_room={self.bootstrap_room}): {self.failure_reason}" + + +# prefill +@dataclasses.dataclass +class TransferKVChunk: + room: int + prefill_kv_indices: npt.NDArray[np.int32] + index_slice: slice + is_last: bool + prefill_aux_index: Optional[int] + + +# decode +@dataclasses.dataclass +class TransferInfo: + room: int + endpoint: str + dst_port: int + mooncake_session_id: str + dst_kv_indices: npt.NDArray[np.int32] + dst_aux_index: int + required_dst_info_num: int + is_dummy: bool + + @classmethod + def from_zmq(cls, msg: List[bytes]): + if msg[4] == b"" and msg[5] == b"": + is_dummy = True + dst_kv_indices = np.array([], dtype=np.int32) + dst_aux_index = None + else: + dst_kv_indices = np.frombuffer(msg[4], dtype=np.int32) + dst_aux_index = int(msg[5].decode("ascii")) + is_dummy = False + return cls( + room=int(msg[0].decode("ascii")), + endpoint=msg[1].decode("ascii"), + dst_port=int(msg[2].decode("ascii")), + mooncake_session_id=msg[3].decode("ascii"), + dst_kv_indices=dst_kv_indices, + dst_aux_index=dst_aux_index, + required_dst_info_num=int(msg[6].decode("ascii")), + is_dummy=is_dummy, + ) + + +# decode +@dataclasses.dataclass +class KVArgsRegisterInfo: + room: str + endpoint: str + dst_port: int + mooncake_session_id: str + dst_kv_ptrs: list[int] + dst_aux_ptrs: list[int] + dst_tp_rank: int + dst_attn_tp_size: int + dst_kv_item_len: int + + @classmethod + def from_zmq(cls, msg: List[bytes]): + return cls( + room=str(msg[0].decode("ascii")), + endpoint=msg[1].decode("ascii"), + dst_port=int(msg[2].decode("ascii")), + mooncake_session_id=msg[3].decode("ascii"), + dst_kv_ptrs=list(struct.unpack(f"{len(msg[4])//8}Q", msg[4])), + dst_aux_ptrs=list(struct.unpack(f"{len(msg[5])//8}Q", msg[5])), + dst_tp_rank=int(msg[6].decode("ascii")), + dst_attn_tp_size=int(msg[7].decode("ascii")), + dst_kv_item_len=int(msg[8].decode("ascii")), + ) + + +class AuxDataCodec: + """Handles serialization and deserialization of auxiliary data buffers""" + + @staticmethod + def serialize_data_from_buffer(src_addr, data_length): + """Serialize data from memory buffer to bytes""" + buffer = (ctypes.c_byte * data_length).from_address(src_addr) + return bytes(buffer) + + @staticmethod + def deserialize_data_to_buffer(kv_args, buffer_index, aux_index, data): + """Deserialize bytes into target memory buffer""" + dst_aux_ptr = kv_args.aux_data_ptrs[buffer_index] + item_len = kv_args.aux_item_lens[buffer_index] + dst_addr = dst_aux_ptr + item_len * aux_index + buffer = (ctypes.c_byte * len(data)).from_address(dst_addr) + buffer[:] = data + return + + +class MooncakeKVManager(BaseKVManager): + AUX_DATA_HEADER = b"AUX_DATA" + + def __init__( + self, + args: KVArgs, + disaggregation_mode: DisaggregationMode, + server_args: ServerArgs, + is_mla_backend: Optional[bool] = False, + ): + self.kv_args = args + self.local_ip = get_local_ip_auto() + self.is_mla_backend = is_mla_backend + self.disaggregation_mode = disaggregation_mode + self.init_engine() + # for p/d multi node infer + self.bootstrap_host = server_args.host + self.bootstrap_port = server_args.disaggregation_bootstrap_port + self.dist_init_addr = server_args.dist_init_addr + self.attn_tp_size = get_attention_tp_size() + self.attn_tp_rank = get_attention_tp_rank() + self.attn_dp_size = get_attention_dp_size() + self.attn_dp_rank = get_attention_dp_rank() + self.system_dp_size = ( + 1 if server_args.enable_dp_attention else server_args.dp_size + ) + self.system_dp_rank = ( + self.kv_args.system_dp_rank if self.kv_args.system_dp_rank else 0 + ) + self.pp_size = server_args.pp_size + self.pp_rank = self.kv_args.pp_rank + self.request_status: Dict[int, KVPoll] = {} + self.rank_port = None + self.server_socket = zmq.Context().socket(zmq.PULL) + if is_valid_ipv6_address(self.local_ip): + self.server_socket.setsockopt(zmq.IPV6, 1) + + self.register_buffer_to_engine() + if self.disaggregation_mode == DisaggregationMode.PREFILL: + self.transfer_infos: Dict[int, Dict[str, TransferInfo]] = {} + self.decode_kv_args_table: Dict[str, KVArgsRegisterInfo] = {} + self.start_prefill_thread() + self._register_to_bootstrap() + self.session_failures = defaultdict(int) + self.failed_sessions = set() + self.session_lock = threading.Lock() + self.pp_group = get_pp_group() + # Determine the number of threads to use for kv sender + cpu_count = os.cpu_count() + transfer_thread_pool_size = get_int_env_var( + "SGLANG_DISAGGREGATION_THREAD_POOL_SIZE", + min(max(4, int(0.75 * cpu_count) // 8), 12), + ) + transfer_queue_size = get_int_env_var("SGLANG_DISAGGREGATION_QUEUE_SIZE", 4) + self.transfer_queues: List[FastQueue] = [ + FastQueue() for _ in range(transfer_queue_size) + ] + assert transfer_thread_pool_size >= transfer_queue_size, ( + f"The environment variable SGLANG_DISAGGREGATION_THREAD_POOL_SIZE={transfer_thread_pool_size} must be " + f"greater than or equal to SGLANG_DISAGGREGATION_QUEUE_SIZE={transfer_queue_size}." + ) + self.executors = [ + concurrent.futures.ThreadPoolExecutor( + transfer_thread_pool_size // transfer_queue_size + ) + for _ in range(transfer_queue_size) + ] + for queue, executor in zip(self.transfer_queues, self.executors): + threading.Thread( + target=self.transfer_worker, args=(queue, executor), daemon=True + ).start() + # If a timeout happens on the prefill side, it means prefill instances + # fail to receive the KV indices from the decode instance of this request. + # These timeout requests should be aborted to release the tree cache. + self.bootstrap_timeout = get_int_env_var( + "SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT", 300 + ) + + self.enable_custom_mem_pool = get_bool_env_var( + "SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false" + ) + elif self.disaggregation_mode == DisaggregationMode.DECODE: + self.heartbeat_failures = {} + self.session_pool = defaultdict(requests.Session) + self.session_pool_lock = threading.Lock() + self.addr_to_rooms_tracker = defaultdict(set) + self.connection_lock = threading.Lock() + self.required_prefill_response_num_table: Dict[int, int] = {} + self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set) + # Heartbeat interval should be at least 2 seconds + self.heartbeat_interval = max( + float(os.getenv("SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL", 5.0)), 2.0 + ) + # Heartbeat failure should be at least 1 + self.max_failures = max( + get_int_env_var("SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE", 2), 1 + ) + self.start_decode_thread() + self.connection_pool: Dict[str, Dict[str, Union[str, int]]] = {} + self.prefill_attn_tp_size_table: Dict[str, int] = {} + self.prefill_dp_size_table: Dict[str, int] = {} + self.prefill_pp_size_table: Dict[str, int] = {} + # If a timeout happens on the decode side, it means decode instances + # fail to receive the KV Cache transfer done signal after bootstrapping. + # These timeout requests should be aborted to release the tree cache. + self.waiting_timeout = get_int_env_var( + "SGLANG_DISAGGREGATION_WAITING_TIMEOUT", 300 + ) + else: + raise ValueError( + f"Unsupported DisaggregationMode: {self.disaggregation_mode}" + ) + + self.failure_records: Dict[int, str] = {} + self.failure_lock = threading.Lock() + + def init_engine(self): + self.engine = MooncakeTransferEngine( + hostname=self.local_ip, + gpu_id=self.kv_args.gpu_id, + ib_device=self.kv_args.ib_device, + ) + + def register_buffer_to_engine(self): + # Batch register KV data buffers + if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens: + self.engine.batch_register( + self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens + ) + + # Batch register auxiliary data buffers + if self.kv_args.aux_data_ptrs and self.kv_args.aux_data_lens: + self.engine.batch_register( + self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens + ) + + @cache + def _connect(self, endpoint: str, is_ipv6: bool = False): + socket = zmq.Context().socket(zmq.PUSH) + if is_ipv6: + socket.setsockopt(zmq.IPV6, 1) + socket.connect(endpoint) + return socket + + def _transfer_data(self, mooncake_session_id, transfer_blocks): + if not transfer_blocks: + return 0 + + src_addrs, dst_addrs, lengths = zip(*transfer_blocks) + return self.engine.batch_transfer_sync( + mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths) + ) + + def send_kvcache( + self, + mooncake_session_id: str, + prefill_kv_indices: npt.NDArray[np.int32], + dst_kv_ptrs: list[int], + dst_kv_indices: npt.NDArray[np.int32], + executor: concurrent.futures.ThreadPoolExecutor, + ): + # Group by indices + prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( + prefill_kv_indices, dst_kv_indices + ) + + layers_params = None + + # pp is not supported on the decode side yet + start_layer = self.kv_args.prefill_start_layer + end_layer = start_layer + len(self.kv_args.kv_data_ptrs) + if self.is_mla_backend: + src_kv_ptrs = self.kv_args.kv_data_ptrs + layers_per_pp_stage = len(src_kv_ptrs) + dst_kv_ptrs = dst_kv_ptrs[start_layer:end_layer] + kv_item_len = self.kv_args.kv_item_lens[0] + layers_params = [ + ( + src_kv_ptrs[layer_id], + dst_kv_ptrs[layer_id], + kv_item_len, + ) + for layer_id in range(layers_per_pp_stage) + ] + else: + num_kv_layers = len(self.kv_args.kv_data_ptrs) // 2 + dst_num_total_layers = num_kv_layers * self.pp_size + src_k_ptrs = self.kv_args.kv_data_ptrs[:num_kv_layers] + src_v_ptrs = self.kv_args.kv_data_ptrs[num_kv_layers:] + layers_per_pp_stage = len(src_k_ptrs) + dst_k_ptrs = dst_kv_ptrs[start_layer:end_layer] + dst_v_ptrs = dst_kv_ptrs[ + dst_num_total_layers + start_layer : dst_num_total_layers + end_layer + ] + kv_item_len = self.kv_args.kv_item_lens[0] + layers_params = [ + ( + src_k_ptrs[layer_id], + dst_k_ptrs[layer_id], + kv_item_len, + ) + for layer_id in range(layers_per_pp_stage) + ] + [ + ( + src_v_ptrs[layer_id], + dst_v_ptrs[layer_id], + kv_item_len, + ) + for layer_id in range(layers_per_pp_stage) + ] + assert layers_params is not None + + def set_transfer_blocks( + src_ptr: int, dst_ptr: int, item_len: int + ) -> List[Tuple[int, int, int]]: + transfer_blocks = [] + for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks): + src_addr = src_ptr + int(prefill_index[0]) * item_len + dst_addr = dst_ptr + int(decode_index[0]) * item_len + length = item_len * len(prefill_index) + transfer_blocks.append((src_addr, dst_addr, length)) + return transfer_blocks + + # Worker function for processing a single layer + def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int: + transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len) + return self._transfer_data(mooncake_session_id, transfer_blocks) + + # Worker function for processing all layers in a batch + def process_layers(layers_params: List[Tuple[int, int, int]]) -> int: + transfer_blocks = [] + for src_ptr, dst_ptr, item_len in layers_params: + transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len)) + return self._transfer_data(mooncake_session_id, transfer_blocks) + + if self.enable_custom_mem_pool: + futures = [ + executor.submit( + process_layer, + src_ptr, + dst_ptr, + item_len, + ) + for (src_ptr, dst_ptr, item_len) in layers_params + ] + for future in concurrent.futures.as_completed(futures): + status = future.result() + if status != 0: + for f in futures: + f.cancel() + return status + else: + # Combining all layers' params in one batch transfer is more efficient + # compared to using multiple threads + return process_layers(layers_params) + + return 0 + + def send_kvcache_slice( + self, + mooncake_session_id: str, + prefill_kv_indices: npt.NDArray[np.int64], + dst_kv_ptrs: list[int], + dst_kv_indices: npt.NDArray[np.int64], + dst_tp_rank: int, + dst_attn_tp_size: int, + dst_kv_item_len: int, + executor: concurrent.futures.ThreadPoolExecutor, + ): + """ + Sends KV cache slices from this Prefill rank to a target Decode rank, + supporting generic M-to-N TP size configurations. + + NOTE: This implementation calls the transfer engine for each token slot within + each page to ensure correctness for any page_size and head-slicing configuration. + This may introduce performance overhead (increased TTFT) for long sequences. + """ + # Extract configuration + local_tp_rank_in_group = self.kv_args.engine_rank % self.attn_tp_size + src_kv_item_len = self.kv_args.kv_item_lens[0] + dst_tp_rank_in_group = dst_tp_rank % dst_attn_tp_size + num_kv_heads = self.kv_args.kv_head_num + num_layers = len(self.kv_args.kv_data_ptrs) + page_size = self.kv_args.page_size + + # Calculate head distribution + src_heads_per_rank = num_kv_heads + dst_heads_per_rank = num_kv_heads * self.attn_tp_size // dst_attn_tp_size + bytes_per_head_slice_to_send = ( + dst_kv_item_len // page_size // dst_heads_per_rank + ) + + # Determine slicing parameters based on TP configuration + if self.attn_tp_size > dst_attn_tp_size: + # Send KVCache from multiple prefill instances to 1 decode instance + src_head_start_offset = 0 + num_heads_to_send = src_heads_per_rank + dst_head_start_offset = local_tp_rank_in_group * src_heads_per_rank + else: + # Send KVCache from 1 prefill instance to multiple decode instances + src_head_start_offset = ( + dst_tp_rank_in_group * dst_heads_per_rank + ) % src_heads_per_rank + num_heads_to_send = dst_heads_per_rank + dst_head_start_offset = 0 + + # pp is not supported on the decode side yet + num_kv_layers = len(self.kv_args.kv_data_ptrs) // 2 + dst_num_total_layers = num_kv_layers * self.pp_size + src_k_ptrs = self.kv_args.kv_data_ptrs[:num_kv_layers] + src_v_ptrs = self.kv_args.kv_data_ptrs[num_kv_layers:] + layers_per_pp_stage = len(src_k_ptrs) + start_layer = self.pp_rank * layers_per_pp_stage + end_layer = start_layer + layers_per_pp_stage + dst_k_ptrs = dst_kv_ptrs[start_layer:end_layer] + dst_v_ptrs = dst_kv_ptrs[ + dst_num_total_layers + start_layer : dst_num_total_layers + end_layer + ] + + # Calculate precise byte offset and length for the sub-slice within the token + src_head_slice_offset = src_head_start_offset * bytes_per_head_slice_to_send + dst_head_slice_offset = dst_head_start_offset * bytes_per_head_slice_to_send + heads_bytes_per_token_to_send = num_heads_to_send * bytes_per_head_slice_to_send + + # Sanity check: The data sub-slice to be sent should fit into the dst buffer. + # This means heads_bytes_per_token_to_send <= (dst_kv_item_len // page_size) + if heads_bytes_per_token_to_send > (dst_kv_item_len // page_size): + logger.error( + f"[{mooncake_session_id}] slice size ({heads_bytes_per_token_to_send}) exceeds " + f"target token slot size ({dst_kv_item_len // page_size})" + ) + return -1 + + layers_params = [ + ( + src_k_ptrs[layer_id], + dst_k_ptrs[layer_id], + src_kv_item_len, + dst_kv_item_len, + src_head_slice_offset, + dst_head_slice_offset, + heads_bytes_per_token_to_send, + ) + for layer_id in range(layers_per_pp_stage) + ] + [ + ( + src_v_ptrs[layer_id], + dst_v_ptrs[layer_id], + src_kv_item_len, + dst_kv_item_len, + src_head_slice_offset, + dst_head_slice_offset, + heads_bytes_per_token_to_send, + ) + for layer_id in range(layers_per_pp_stage) + ] + + def process_layer_tp_aware(layer_params): + ( + src_ptr, + dst_ptr, + src_item_len, + dst_item_len, + src_head_slice_offset, + dst_head_slice_offset, + heads_bytes_per_token_to_send, + ) = layer_params + src_addr_list = [] + dst_addr_list = [] + length_list = [] + + # Calculate strides for a single token slot + bytes_per_token_on_prefill = src_item_len // page_size + bytes_per_token_on_decode = dst_item_len // page_size + + for i in range(len(prefill_kv_indices)): + prefill_page_idx = int(prefill_kv_indices[i]) + decode_page_idx = int(dst_kv_indices[i]) + + # Get the starting addresses for the current src and dst pages + src_page_start_addr = src_ptr + prefill_page_idx * src_item_len + dst_page_start_addr = dst_ptr + decode_page_idx * dst_item_len + + # Iterate through each valid token slot within the current page + for token_slot_in_page in range(page_size): + # Calculate the start address of the current token slot + src_token_slot_start_addr = ( + src_page_start_addr + + token_slot_in_page * bytes_per_token_on_prefill + ) + dst_token_slot_start_addr = ( + dst_page_start_addr + + token_slot_in_page * bytes_per_token_on_decode + ) + + # Calculate final src and dst addresses by applying head-slice offsets + src_slice_addr = src_token_slot_start_addr + src_head_slice_offset + dst_slice_addr = dst_token_slot_start_addr + dst_head_slice_offset + + src_addr_list.append(src_slice_addr) + dst_addr_list.append(dst_slice_addr) + length_list.append(heads_bytes_per_token_to_send) + + return self.engine.batch_transfer_sync( + mooncake_session_id, src_addr_list, dst_addr_list, length_list + ) + + futures = [ + executor.submit( + process_layer_tp_aware, + layer_params, + ) + for layer_params in layers_params + ] + + for future in concurrent.futures.as_completed(futures): + status = future.result() + if status != 0: + for f in futures: + f.cancel() + return status + + return 0 + + def send_aux( + self, + req: TransferInfo, + prefill_aux_index: int, + dst_aux_ptrs: list[int], + ): + # TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free + if self.enable_custom_mem_pool: + return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs) + + transfer_blocks = [] + prefill_aux_ptrs = self.kv_args.aux_data_ptrs + prefill_aux_item_lens = self.kv_args.aux_item_lens + + for i, dst_aux_ptr in enumerate(dst_aux_ptrs): + length = prefill_aux_item_lens[i] + src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index + dst_addr = dst_aux_ptrs[i] + length * req.dst_aux_index + transfer_blocks.append((src_addr, dst_addr, length)) + + return self._transfer_data(req.mooncake_session_id, transfer_blocks) + + def send_aux_tcp( + self, + req: TransferInfo, + prefill_aux_index: int, + dst_aux_ptrs: list[int], + ): + prefill_aux_ptrs = self.kv_args.aux_data_ptrs + prefill_aux_item_lens = self.kv_args.aux_item_lens + + for i in range(len(prefill_aux_ptrs)): + length = prefill_aux_item_lens[i] + src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index + data = AuxDataCodec.serialize_data_from_buffer(src_addr, length) + + self.send_aux_data_to_endpoint( + remote=req.endpoint, + dst_port=req.dst_port, + room=req.room, + buffer_index=i, + aux_index=req.dst_aux_index, + data=data, + ) + + return 0 + + def send_aux_data_to_endpoint( + self, + remote: str, + dst_port: int, + room: int, + buffer_index: int, + aux_index: int, + data: bytes, + ): + socket = self._connect( + format_tcp_address(remote, dst_port), is_ipv6=is_valid_ipv6_address(remote) + ) + + socket.send_multipart( + [ + MooncakeKVManager.AUX_DATA_HEADER, + str(room).encode("ascii"), + str(buffer_index).encode("ascii"), + str(aux_index).encode("ascii"), + struct.pack(">I", len(data)), + data, + ] + ) + + def sync_status_to_decode_endpoint( + self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int + ): + self._connect( + format_tcp_address(remote, dst_port), is_ipv6=is_valid_ipv6_address(remote) + ).send_multipart( + [ + str(room).encode("ascii"), + str(status).encode("ascii"), + str(prefill_rank).encode("ascii"), + ] + ) + + def transfer_worker( + self, queue: FastQueue, executor: concurrent.futures.ThreadPoolExecutor + ): + while True: + try: + kv_chunk: TransferKVChunk = queue.get() + reqs_to_be_processed = ( + self.transfer_infos[kv_chunk.room].values() + if kv_chunk.room in self.transfer_infos + else [] + ) + polls = [] + dst_ranks_infos = [] + local_rank = self.attn_tp_rank * self.pp_size + self.pp_rank + for req in reqs_to_be_processed: + if not req.is_dummy: + # Early exit if the request has failed + with self.session_lock: + if req.mooncake_session_id in self.failed_sessions: + self.record_failure( + kv_chunk.room, + f"Decode instance could be dead, remote mooncake session {req.mooncake_session_id} is not alive", + ) + self.update_status(kv_chunk.room, KVPoll.Failed) + self.sync_status_to_decode_endpoint( + req.endpoint, + req.dst_port, + req.room, + KVPoll.Failed, + local_rank, + ) + break + + chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice] + + # NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices + # is mismatched with the dst_kv_indices when page size > 1, this should never happen. + if len(chunked_dst_kv_indice) < len( + kv_chunk.prefill_kv_indices + ): + logger.warning( + f"len(chunked_dst_kv_indice) = {len(chunked_dst_kv_indice)}, len(kv_chunk.prefill_kv_indices) = {len(kv_chunk.prefill_kv_indices)}" + ) + kv_chunk.prefill_kv_indices = kv_chunk.prefill_kv_indices[ + : len(chunked_dst_kv_indice) + ] + + target_rank_registration_info: KVArgsRegisterInfo = ( + self.decode_kv_args_table[req.mooncake_session_id] + ) + if self.is_mla_backend or ( + self.attn_tp_size + == target_rank_registration_info.dst_attn_tp_size + ): + ret = self.send_kvcache( + req.mooncake_session_id, + kv_chunk.prefill_kv_indices, + target_rank_registration_info.dst_kv_ptrs, + chunked_dst_kv_indice, + executor, + ) + else: + ret = self.send_kvcache_slice( + req.mooncake_session_id, + kv_chunk.prefill_kv_indices, + target_rank_registration_info.dst_kv_ptrs, + chunked_dst_kv_indice, + target_rank_registration_info.dst_tp_rank, + target_rank_registration_info.dst_attn_tp_size, + target_rank_registration_info.dst_kv_item_len, + executor, + ) + if ret != 0: + with self.session_lock: + self.session_failures[req.mooncake_session_id] += 1 + # Failures should never happen if the session is not dead, if the session fails once, mark it as failed + if self.session_failures[req.mooncake_session_id] >= 1: + self.failed_sessions.add(req.mooncake_session_id) + logger.error( + f"Session {req.mooncake_session_id} failed." + ) + self.record_failure( + kv_chunk.room, + f"Failed to send kv chunk of {kv_chunk.room} to {req.endpoint}:{req.dst_port}", + ) + self.update_status(kv_chunk.room, KVPoll.Failed) + self.sync_status_to_decode_endpoint( + req.endpoint, + req.dst_port, + req.room, + KVPoll.Failed, + local_rank, + ) + break + + if kv_chunk.is_last: + if self.pp_group.is_last_rank: + # Only the last chunk we need to send the aux data + ret = self.send_aux( + req, + kv_chunk.prefill_aux_index, + target_rank_registration_info.dst_aux_ptrs, + ) + polls.append(True if ret == 0 else False) + dst_ranks_infos.append( + (req.endpoint, req.dst_port, req.room) + ) + + # Only sync status when all the dst ranks have received the kvcache + if len(polls) == req.required_dst_info_num: + status = KVPoll.Success if all(polls) else KVPoll.Failed + self.update_status(req.room, status) + for endpoint, dst_port, room in dst_ranks_infos: + self.sync_status_to_decode_endpoint( + endpoint, dst_port, room, status, local_rank + ) + else: + # Dummy request means the decode instance is not used, so its status can be marked as success directly + # Dummy request does not need to sync status to decode endpoint + if kv_chunk.is_last and req.room in self.request_status: + self.update_status(req.room, KVPoll.Success) + + if ( + kv_chunk.room not in self.request_status + or self.check_status(kv_chunk.room) == KVPoll.Success + ): + if kv_chunk.room in self.transfer_infos: + self.transfer_infos.pop(kv_chunk.room) + + except Exception as e: + # NOTE(shangming): Remove this when we make sure the transfer thread is bug-free + raise RuntimeError( + f"Transfer thread failed because of {e}. Prefill instance with bootstrap_port={self.bootstrap_port} is dead." + ) + + def _bind_server_socket(self): + self.server_socket.bind(format_tcp_address(self.local_ip, self.rank_port)) + + def start_prefill_thread(self): + self.rank_port = get_free_port() + self._bind_server_socket() + + def bootstrap_thread(): + """This thread recvs pre-alloc notification from the decode engine""" + # KVPoll.Bootstrapping -> KVPoll.WaitingForInput + while True: + waiting_req_bytes = self.server_socket.recv_multipart() + room = waiting_req_bytes[0].decode("ascii") + mooncake_session_id = waiting_req_bytes[3].decode("ascii") + if room == "None": + self.decode_kv_args_table[mooncake_session_id] = ( + KVArgsRegisterInfo.from_zmq(waiting_req_bytes) + ) + with self.session_lock: + if mooncake_session_id in self.failed_sessions: + self.failed_sessions.remove(mooncake_session_id) + if mooncake_session_id in self.session_failures: + del self.session_failures[mooncake_session_id] + logger.debug( + f"Register KVArgs from {mooncake_session_id} successfully" + ) + continue + else: + required_dst_info_num = int(waiting_req_bytes[6].decode("ascii")) + room = int(room) + if room not in self.transfer_infos: + self.transfer_infos[room] = {} + + self.transfer_infos[room][mooncake_session_id] = ( + TransferInfo.from_zmq(waiting_req_bytes) + ) + # NOTE: after bootstrapping we can mark the req as waiting for input + if len(self.transfer_infos[room]) == required_dst_info_num: + self.update_status(room, KVPoll.WaitingForInput) + + threading.Thread(target=bootstrap_thread).start() + + def _handle_aux_data(self, msg: List[bytes]): + """Handle AUX_DATA messages received by the decode thread.""" + room = int(msg[1].decode("ascii")) + buffer_index = int(msg[2].decode("ascii")) + aux_index = int(msg[3].decode("ascii")) + data_length = struct.unpack(">I", msg[4])[0] + data = msg[5] + + if len(data) != data_length: + logger.error(f"AUX_DATA length mismatch for bootstrap_room {room}") + return + + AuxDataCodec.deserialize_data_to_buffer( + self.kv_args, buffer_index, aux_index, data + ) + + logger.debug( + f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}" + ) + + def start_decode_thread(self): + self.rank_port = get_free_port() + self._bind_server_socket() + + def decode_thread(): + while True: + msg = self.server_socket.recv_multipart() + if msg[0] == MooncakeKVManager.AUX_DATA_HEADER: + self._handle_aux_data(msg) + continue + + (bootstrap_room, status, prefill_rank) = msg + status = int(status.decode("ascii")) + bootstrap_room = int(bootstrap_room.decode("ascii")) + prefill_rank = int(prefill_rank.decode("ascii")) + + if status == KVPoll.Success: + if bootstrap_room in self.request_status: + self.prefill_response_tracker[bootstrap_room].add(prefill_rank) + expected_response_num = ( + self.required_prefill_response_num_table[bootstrap_room] + ) + arrived_response_num = len( + self.prefill_response_tracker[bootstrap_room] + ) + if arrived_response_num == expected_response_num: + self.update_status(bootstrap_room, KVPoll.Success) + elif status == KVPoll.Failed: + self.record_failure( + bootstrap_room, + f"Failed to get kvcache from prefill instance, it might be dead", + ) + self.update_status(bootstrap_room, status) + + def heartbeat_checker(): + while True: + time.sleep(self.heartbeat_interval) + with self.connection_lock: + addresses = list(self.prefill_dp_size_table.keys()) + + for bootstrap_addr in addresses: + session = None + try: + with self.session_pool_lock: + session = self.session_pool[bootstrap_addr] + response = session.get( + f"http://{bootstrap_addr}/health", + timeout=(2, 3), + headers={"Connection": "keep-alive"}, + ) + if response.status_code == 200: + self.heartbeat_failures[bootstrap_addr] = 0 + + current_rooms = self.addr_to_rooms_tracker[ + bootstrap_addr + ].copy() + + for bootstrap_room in current_rooms: + # Remove KVPoll.Success requests from the tracker + if bootstrap_room not in self.request_status: + self.addr_to_rooms_tracker[bootstrap_addr].discard( + bootstrap_room + ) + else: + logger.info( + f"Attempting to reconnect to {bootstrap_addr}..." + ) + self.heartbeat_failures[bootstrap_addr] = ( + self.heartbeat_failures.get(bootstrap_addr, 0) + 1 + ) + with self.session_pool_lock: + if bootstrap_addr in self.session_pool: + del self.session_pool[bootstrap_addr] + except Exception: + logger.info(f"Attempting to reconnect to {bootstrap_addr}...") + self.heartbeat_failures[bootstrap_addr] = ( + self.heartbeat_failures.get(bootstrap_addr, 0) + 1 + ) + + if ( + self.heartbeat_failures.get(bootstrap_addr, 0) + >= self.max_failures + ): + self._handle_node_failure(bootstrap_addr) + with self.session_pool_lock: + if bootstrap_addr in self.session_pool: + del self.session_pool[bootstrap_addr] + + threading.Thread(target=decode_thread).start() + threading.Thread(target=heartbeat_checker).start() + + def add_transfer_request( + self, + bootstrap_room: int, + kv_indices: npt.NDArray[np.int32], + index_slice: slice, + is_last: bool, + aux_index: Optional[int] = None, + ): + assert self.disaggregation_mode == DisaggregationMode.PREFILL + assert not is_last or (is_last and aux_index is not None) + + if ( + bootstrap_room not in self.request_status + or self.check_status(bootstrap_room) == KVPoll.Failed + ): + logger.debug( + "Request with bootstrap_room=%s already failed", bootstrap_room + ) + return + + if bootstrap_room not in self.transfer_infos: + # This means that the current rank is a dummy rank for this request, + # and it has already been marked as success, so there is no need to + # add further chunks into the transfer queue. + return + + # NOTE(shangming): sharding according to the dst_infos to make sure + # requests with the same dst_sessions will be added into the same + # queue, which enables early abort with failed sessions. + dst_infos = self.transfer_infos[bootstrap_room].keys() + session_port_sum = sum(int(session.rsplit(":", 1)[1]) for session in dst_infos) + shard_idx = session_port_sum % len(self.transfer_queues) + + self.transfer_queues[shard_idx].put( + TransferKVChunk( + room=bootstrap_room, + prefill_kv_indices=kv_indices, + index_slice=index_slice, + is_last=is_last, + prefill_aux_index=aux_index, + ) + ) + + def check_status(self, bootstrap_room: int): + return self.request_status[bootstrap_room] + + def update_status(self, bootstrap_room: int, status: KVPoll): + if bootstrap_room not in self.request_status: + self.request_status[bootstrap_room] = status + else: + # NOTE: status is only allowed to be incremented unless it is KVPoll.Failed + if status == KVPoll.Failed: + self.request_status[bootstrap_room] = KVPoll.Failed + else: + self.request_status[bootstrap_room] = max( + self.request_status[bootstrap_room], status + ) + + def record_failure(self, bootstrap_room: int, failure_reason: str): + with self.failure_lock: + self.failure_records[bootstrap_room] = failure_reason + + def get_session_id(self): + return self.engine.get_session_id() + + def _register_to_bootstrap(self): + """Register KVSender to bootstrap server via HTTP POST.""" + if self.dist_init_addr: + # multi node case: bootstrap server's host is dist_init_addr + if self.dist_init_addr.startswith("["): # [ipv6]:port or [ipv6] + if self.dist_init_addr.endswith("]"): + host = self.dist_init_addr + else: + host, _ = self.dist_init_addr.rsplit(":", 1) + else: + host = socket.gethostbyname(self.dist_init_addr.rsplit(":", 1)[0]) + else: + # single node case: bootstrap server's host is same as http server's host + host = self.bootstrap_host + host = maybe_wrap_ipv6_address(host) + + bootstrap_server_url = f"{host}:{self.bootstrap_port}" + url = f"http://{bootstrap_server_url}/route" + payload = { + "role": "Prefill", + "attn_tp_size": self.attn_tp_size, + "attn_tp_rank": self.attn_tp_rank, + "attn_dp_size": self.attn_dp_size, + "attn_dp_rank": self.attn_dp_rank, + "pp_size": self.pp_size, + "pp_rank": self.pp_rank, + "system_dp_size": self.system_dp_size, + "system_dp_rank": self.system_dp_rank, + "rank_ip": self.local_ip, + "rank_port": self.rank_port, + } + + try: + response = requests.put(url, json=payload, timeout=5) + if response.status_code == 200: + logger.debug("Prefill successfully registered to bootstrap server.") + else: + logger.error( + f"Prefill instance failed to connect to bootstrap server: {response.status_code}, {response.text}" + ) + except Exception as e: + logger.error( + f"Prefill instance failed to register to bootstrap server: {e}" + ) + + def _handle_node_failure(self, failed_bootstrap_addr): + with self.connection_lock: + keys_to_remove = [ + k for k in self.connection_pool if k.startswith(failed_bootstrap_addr) + ] + for k in keys_to_remove: + del self.connection_pool[k] + if failed_bootstrap_addr in self.prefill_attn_tp_size_table: + del self.prefill_attn_tp_size_table[failed_bootstrap_addr] + if failed_bootstrap_addr in self.prefill_dp_size_table: + del self.prefill_dp_size_table[failed_bootstrap_addr] + if failed_bootstrap_addr in self.prefill_pp_size_table: + del self.prefill_pp_size_table[failed_bootstrap_addr] + + possible_affected_rooms = self.addr_to_rooms_tracker.get( + failed_bootstrap_addr, [] + ) + if failed_bootstrap_addr in self.addr_to_rooms_tracker: + del self.addr_to_rooms_tracker[failed_bootstrap_addr] + + # Report the requests associated with the failed bootstrap addr and mark their status as KVPoll.Failed + affected_rooms = [] + for room in possible_affected_rooms: + if ( + room in self.request_status + and self.check_status(room) != KVPoll.Success + ): + self.record_failure( + room, + f"Losing connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr})", + ) + self.update_status(room, KVPoll.Failed) + affected_rooms.append(room) + logger.error( + f"Losing connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr}), {len(affected_rooms)} requests affected" + ) + + +class MooncakeKVSender(BaseKVSender): + + def __init__( + self, + mgr: MooncakeKVManager, + bootstrap_addr: str, + bootstrap_room: int, + dest_tp_ranks: List[int], + pp_rank: int, + ): + self.kv_mgr = mgr + self.bootstrap_room = bootstrap_room + self.kv_mgr.update_status(bootstrap_room, KVPoll.Bootstrapping) + self.aux_index = None + self.bootstrap_server_url = bootstrap_addr + self.conclude_state = None + self.init_time = time.time() + # inner state + self.curr_idx = 0 + + def init(self, num_kv_indices: int, aux_index: Optional[int] = None): + self.num_kv_indices = num_kv_indices + self.aux_index = aux_index + + def send( + self, + kv_indices: npt.NDArray[np.int32], + ): + index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices)) + self.curr_idx += len(kv_indices) + is_last = self.curr_idx == self.num_kv_indices + + if not is_last: + self.kv_mgr.add_transfer_request( + self.bootstrap_room, + kv_indices, + index_slice, + False, + ) + else: + self.kv_mgr.add_transfer_request( + self.bootstrap_room, + kv_indices, + index_slice, + True, + aux_index=self.aux_index, + ) + + def poll(self) -> KVPoll: + if self.conclude_state is None: + status = self.kv_mgr.check_status(self.bootstrap_room) + if status in (KVPoll.Success, KVPoll.Failed): + self.conclude_state = status + elif status == KVPoll.Bootstrapping: + if self.init_time is not None: + now = time.time() + elapsed = now - self.init_time + if elapsed >= self.kv_mgr.bootstrap_timeout: + logger.warning_once( + "Some requests timed out when bootstrapping, " + "which means prefill instances fail to receive the KV indices from the decode instance of this request. " + "If a greater mean TTFT is acceptable, you can 'export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600' (10 minutes) to relax the timeout condition. " + ) + self.kv_mgr.record_failure( + self.bootstrap_room, + f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.Bootstrapping", + ) + self.conclude_state = KVPoll.Failed + return KVPoll.Failed + + return status + else: + return self.conclude_state + + def clear(self) -> None: + if self.bootstrap_room in self.kv_mgr.request_status: + self.kv_mgr.request_status.pop(self.bootstrap_room) + + def failure_exception(self): + # Explicitly set the status to failure since this request has failed in another rank + if self.conclude_state is None: + self.conclude_state = KVPoll.Failed + + self.clear() + + with self.kv_mgr.failure_lock: + failure_reason = self.kv_mgr.failure_records.pop( + self.bootstrap_room, "Failed due to an unknown reason from another rank" + ) + raise KVTransferError(self.bootstrap_room, failure_reason) + + def abort(self): + self.kv_mgr.record_failure( + self.bootstrap_room, + "Aborted by AbortReq.", + ) + # Explicitly set the status to failure since this request has been aborted + self.conclude_state = KVPoll.Failed + + +class MooncakeKVReceiver(BaseKVReceiver): + _ctx = zmq.Context() + _socket_cache = {} + _socket_locks = {} + _global_lock = threading.Lock() + + def __init__( + self, + mgr: MooncakeKVManager, + bootstrap_addr: str, + bootstrap_room: Optional[int] = None, + prefill_dp_rank: Optional[int] = None, + ): + self.bootstrap_room = bootstrap_room + self.bootstrap_addr = bootstrap_addr + self.kv_mgr = mgr + self.session_id = self.kv_mgr.get_session_id() + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) + self.conclude_state = None + self.init_time = None + + if self.bootstrap_addr not in self.kv_mgr.prefill_dp_size_table: + ( + self.prefill_attn_tp_size, + self.prefill_dp_size, + self.prefill_pp_size, + ) = self._get_prefill_parallel_info_from_server() + if ( + self.prefill_attn_tp_size is None + or self.prefill_dp_size is None + or self.prefill_pp_size is None + ): + self.kv_mgr.record_failure( + self.bootstrap_room, + f"Could not fetch prefill parallel info from bootstrap_addr: {self.bootstrap_addr}", + ) + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + return + else: + logger.debug( + f"Fetch prefill parallel info from [{self.bootstrap_addr}]: DP size:{self.prefill_dp_size}, TP size:{self.prefill_attn_tp_size} PP size:{self.prefill_pp_size}" + ) + self.kv_mgr.prefill_attn_tp_size_table[self.bootstrap_addr] = ( + self.prefill_attn_tp_size + ) + self.kv_mgr.prefill_dp_size_table[self.bootstrap_addr] = ( + self.prefill_dp_size + ) + self.kv_mgr.prefill_pp_size_table[self.bootstrap_addr] = ( + self.prefill_pp_size + ) + else: + self.prefill_attn_tp_size = self.kv_mgr.prefill_attn_tp_size_table[ + self.bootstrap_addr + ] + self.prefill_dp_size = self.kv_mgr.prefill_dp_size_table[ + self.bootstrap_addr + ] + self.prefill_pp_size = self.kv_mgr.prefill_pp_size_table[ + self.bootstrap_addr + ] + + # Currently, we don't allow prefill instance and decode instance to + # have different TP sizes per DP rank, except for models using MLA. + if self.kv_mgr.attn_tp_size == self.prefill_attn_tp_size: + self.target_tp_rank = ( + self.kv_mgr.kv_args.engine_rank % self.kv_mgr.attn_tp_size + ) + self.required_dst_info_num = 1 + self.required_prefill_response_num = 1 * ( + self.prefill_pp_size // self.kv_mgr.pp_size + ) + self.target_tp_ranks = [self.target_tp_rank] + elif self.kv_mgr.attn_tp_size > self.prefill_attn_tp_size: + if not self.kv_mgr.is_mla_backend: + logger.warning_once( + "Performance is NOT guaranteed when using different TP sizes for non-MLA models. " + ) + self.target_tp_rank = ( + self.kv_mgr.kv_args.engine_rank % self.kv_mgr.attn_tp_size + ) // (self.kv_mgr.attn_tp_size // self.prefill_attn_tp_size) + self.required_dst_info_num = ( + self.kv_mgr.attn_tp_size // self.prefill_attn_tp_size + ) + self.required_prefill_response_num = 1 * ( + self.prefill_pp_size // self.kv_mgr.pp_size + ) + self.target_tp_ranks = [self.target_tp_rank] + else: + if not self.kv_mgr.is_mla_backend: + logger.warning_once( + "Performance is NOT guaranteed when using different TP sizes for non-MLA models. " + ) + # For non-MLA models, one decode rank needs to retrieve KVCache from multiple prefill ranks for non MLA models; + self.target_tp_ranks = [ + rank + for rank in range( + (self.kv_mgr.kv_args.engine_rank % self.kv_mgr.attn_tp_size) + * (self.prefill_attn_tp_size // self.kv_mgr.attn_tp_size), + (self.kv_mgr.kv_args.engine_rank % self.kv_mgr.attn_tp_size + 1) + * (self.prefill_attn_tp_size // self.kv_mgr.attn_tp_size), + ) + ] + + # For MLA models, we can retrieve KVCache from only one prefill rank, but we still need to maintain + # multiple connections in the connection pool and have to send dummy requests to other prefill ranks, + # or the KVPoll will never be set correctly + self.target_tp_rank = self.target_tp_ranks[0] + self.required_dst_info_num = 1 + if self.kv_mgr.is_mla_backend: + self.required_prefill_response_num = ( + self.prefill_pp_size // self.kv_mgr.pp_size + ) + else: + self.required_prefill_response_num = ( + self.prefill_attn_tp_size // self.kv_mgr.attn_tp_size + ) * (self.prefill_pp_size // self.kv_mgr.pp_size) + + if prefill_dp_rank is not None: + logger.debug(f"Targeting DP rank: {prefill_dp_rank}") + self.prefill_dp_rank = prefill_dp_rank + else: + self.prefill_dp_rank = bootstrap_room % self.prefill_dp_size + + # FIXME: alias here: target_dp_group -> prefill_dp_rank + self.target_dp_group = self.prefill_dp_rank + + self.kv_mgr.required_prefill_response_num_table[self.bootstrap_room] = ( + self.required_prefill_response_num + ) + # NOTE: key distinguished by bootstrap_addr, target_dp_group, and target_tp_rank + bootstrap_key = ( + f"{self.bootstrap_addr}_{self.target_dp_group}_{self.target_tp_rank}" + ) + + if bootstrap_key not in self.kv_mgr.connection_pool: + bootstrap_infos = [] + for target_tp_rank in self.target_tp_ranks: + for target_pp_rank in range(self.prefill_pp_size): + bootstrap_info = self._get_bootstrap_info_from_server( + target_tp_rank, self.target_dp_group, target_pp_rank + ) + if bootstrap_info is not None: + if self.kv_mgr.is_mla_backend: + # For MLA: target_tp_rank is the selected real rank, others are dummy ranks + bootstrap_info["is_dummy"] = not bool( + target_tp_rank == self.target_tp_rank + or self.target_tp_rank is None + ) + else: + # For non-MLA: all target_tp_ranks are selected real ranks + bootstrap_info["is_dummy"] = False + logger.debug( + f"Fetched bootstrap info: {bootstrap_info} for DP {self.target_dp_group} TP {target_tp_rank} PP {target_pp_rank}" + ) + bootstrap_infos.append(bootstrap_info) + else: + self.kv_mgr.record_failure( + self.bootstrap_room, + f"Could not fetch bootstrap info for engine rank: {self.kv_mgr.kv_args.engine_rank} and target_dp_group: {self.target_dp_group} and target_pp_rank {target_pp_rank}", + ) + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + return + + self.bootstrap_infos = bootstrap_infos + self.kv_mgr.connection_pool[bootstrap_key] = self.bootstrap_infos + + # Register kv_args only once to prefill KVManager according to the info fetched from the bootstrap server + self._register_kv_args() + else: + self.bootstrap_infos = self.kv_mgr.connection_pool[bootstrap_key] + + assert len(self.bootstrap_infos) > 0 + self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput) + + def _get_bootstrap_info_from_server( + self, engine_rank, target_dp_group, target_pp_rank + ): + """Fetch the bootstrap info from the bootstrap server.""" + try: + url = f"http://{self.bootstrap_addr}/route?engine_rank={engine_rank}&target_dp_group={target_dp_group}&target_pp_rank={target_pp_rank}" + response = requests.get(url, timeout=5) + if response.status_code == 200: + bootstrap_info = response.json() + return bootstrap_info + else: + logger.error( + f"Failed to get prefill server info: {response.status_code}, {response.text}" + ) + return None + except Exception as e: + logger.error(f"Error fetching prefill info from bootstrap: {e}") + return None + + def _get_prefill_parallel_info_from_server( + self, + ) -> Tuple[Optional[int], Optional[int], Optional[int]]: + """Fetch the prefill parallel info from the bootstrap server.""" + try: + url = f"http://{self.bootstrap_addr}/route?engine_rank={-1}&target_dp_group={-1}&target_pp_rank={-1}" + response = requests.get(url) + if response.status_code == 200: + prefill_parallel_info = response.json() + return ( + int(prefill_parallel_info["prefill_attn_tp_size"]), + int(prefill_parallel_info["prefill_dp_size"]), + int(prefill_parallel_info["prefill_pp_size"]), + ) + else: + logger.error( + f"Failed to get prefill parallel info: {response.status_code}, {response.text}" + ) + return None, None, None + except Exception as e: + logger.error(f"Error fetching prefill parallel info from bootstrap: {e}") + return None, None, None + + def _register_kv_args(self): + for bootstrap_info in self.bootstrap_infos: + packed_kv_data_ptrs = b"".join( + struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs + ) + packed_aux_data_ptrs = b"".join( + struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs + ) + # Note(shangming): No need to add pp rank here since pp is not supported on the decode side yet + tp_rank = self.kv_mgr.kv_args.engine_rank + kv_item_len = self.kv_mgr.kv_args.kv_item_lens[0] + dst_tp_rank = str(tp_rank).encode("ascii") + dst_attn_tp_size = str(self.kv_mgr.attn_tp_size).encode("ascii") + dst_kv_item_len = str(kv_item_len).encode("ascii") + + sock, lock = self._connect_to_bootstrap_server(bootstrap_info) + with lock: + sock.send_multipart( + [ + "None".encode("ascii"), + self.kv_mgr.local_ip.encode("ascii"), + str(self.kv_mgr.rank_port).encode("ascii"), + self.session_id.encode("ascii"), + packed_kv_data_ptrs, + packed_aux_data_ptrs, + dst_tp_rank, + dst_attn_tp_size, + dst_kv_item_len, + ] + ) + + @classmethod + def _connect(cls, endpoint: str, is_ipv6: bool = False): + with cls._global_lock: + if endpoint not in cls._socket_cache: + sock = cls._ctx.socket(zmq.PUSH) + if is_ipv6: + sock.setsockopt(zmq.IPV6, 1) + sock.connect(endpoint) + cls._socket_cache[endpoint] = sock + cls._socket_locks[endpoint] = threading.Lock() + return cls._socket_cache[endpoint], cls._socket_locks[endpoint] + + @classmethod + def _connect_to_bootstrap_server(cls, bootstrap_info: dict): + ip_address = bootstrap_info["rank_ip"] + port = bootstrap_info["rank_port"] + is_ipv6_address = is_valid_ipv6_address(ip_address) + sock, lock = cls._connect( + format_tcp_address(ip_address, port), is_ipv6=is_ipv6_address + ) + return sock, lock + + def init(self, kv_indices: npt.NDArray[np.int32], aux_index: Optional[int] = None): + for bootstrap_info in self.bootstrap_infos: + sock, lock = self._connect_to_bootstrap_server(bootstrap_info) + is_dummy = bootstrap_info["is_dummy"] + + with lock: + sock.send_multipart( + [ + str(self.bootstrap_room).encode("ascii"), + self.kv_mgr.local_ip.encode("ascii"), + str(self.kv_mgr.rank_port).encode("ascii"), + self.session_id.encode("ascii"), + kv_indices.tobytes() if not is_dummy else b"", + str(aux_index).encode("ascii") if not is_dummy else b"", + str(self.required_dst_info_num).encode("ascii"), + ] + ) + self.init_time = time.time() + + def poll(self) -> KVPoll: + if self.conclude_state is None: + status = self.kv_mgr.check_status(self.bootstrap_room) + if status in (KVPoll.Success, KVPoll.Failed): + self.conclude_state = status + elif status == KVPoll.WaitingForInput: + if self.init_time is not None: + now = time.time() + elapsed = now - self.init_time + if elapsed >= self.kv_mgr.waiting_timeout: + logger.warning_once( + "Some requests fail to receive KV Cache transfer done signal after bootstrapping. " + "If a greater mean TTFT is acceptable, you can 'export SGLANG_DISAGGREGATION_WAITING_TIMEOUT=600' (10 minutes) to relax the timeout condition. " + ) + self.kv_mgr.record_failure( + self.bootstrap_room, + f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.WaitingForInput", + ) + self.conclude_state = KVPoll.Failed + return KVPoll.Failed + + return status + + else: + return self.conclude_state + + def clear(self) -> None: + if self.bootstrap_room in self.kv_mgr.request_status: + self.kv_mgr.request_status.pop(self.bootstrap_room) + + if self.bootstrap_room in self.kv_mgr.required_prefill_response_num_table: + self.kv_mgr.required_prefill_response_num_table.pop(self.bootstrap_room) + + if self.bootstrap_room in self.kv_mgr.prefill_response_tracker: + self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room) + + def failure_exception(self): + # Explicitly set the status to failure since this request has failed in another rank + if self.conclude_state is None: + self.conclude_state = KVPoll.Failed + + self.clear() + + with self.kv_mgr.failure_lock: + failure_reason = self.kv_mgr.failure_records.pop( + self.bootstrap_room, "Failed due to an unknown reason from another rank" + ) + raise KVTransferError(self.bootstrap_room, failure_reason) + + def abort(self): + self.kv_mgr.record_failure( + self.bootstrap_room, + "Aborted by AbortReq.", + ) + # Explicitly set the status to failure since this request has been aborted + self.conclude_state = KVPoll.Failed + + +class MooncakeKVBootstrapServer(BaseKVBootstrapServer): + def __init__(self, host: str, port: int): + self.host = host + self.port = port + self.app = web.Application() + self.store = dict() + self.lock = asyncio.Lock() + self._setup_routes() + self.pp_size = None + self.attn_tp_size = None + self.dp_size = None + self.prefill_port_table: Dict[ + int, Dict[int, Dict[int, Dict[str, Union[str, int]]]] + ] = {} + + # Start bootstrap server + self.thread = threading.Thread(target=self._run_server, daemon=True) + self.run() + + def run(self): + self.thread.start() + + def _setup_routes(self): + self.app.router.add_route("*", "/route", self._handle_route) + self.app.router.add_get("/health", self._handle_health_check) + + async def _handle_health_check(self, request): + return web.Response(text="OK", status=200) + + async def _handle_route(self, request: web.Request): + method = request.method + if method == "PUT": + return await self._handle_route_put(request) + elif method == "GET": + return await self._handle_route_get(request) + else: + return web.Response( + text="Method not allowed", status=405, content_type="application/json" + ) + + async def _handle_route_put(self, request: web.Request): + data = await request.json() + role = data["role"] + attn_tp_size = data["attn_tp_size"] + attn_tp_rank = data["attn_tp_rank"] + attn_dp_size = data["attn_dp_size"] + attn_dp_rank = data["attn_dp_rank"] + pp_size = data["pp_size"] + pp_rank = data["pp_rank"] + system_dp_size = data["system_dp_size"] + system_dp_rank = data["system_dp_rank"] + rank_ip = data["rank_ip"] + rank_port = int(data["rank_port"]) + + if self.attn_tp_size is None: + self.attn_tp_size = attn_tp_size + + if self.dp_size is None: + self.dp_size = attn_dp_size if system_dp_size == 1 else system_dp_size + + if self.pp_size is None: + self.pp_size = pp_size + + if role == "Prefill": + if system_dp_size == 1: + dp_group = attn_dp_rank + else: + dp_group = system_dp_rank + + # Add lock to make sure thread-safe + async with self.lock: + if dp_group not in self.prefill_port_table: + self.prefill_port_table[dp_group] = {} + if attn_tp_rank not in self.prefill_port_table[dp_group]: + self.prefill_port_table[dp_group][attn_tp_rank] = {} + + self.prefill_port_table[dp_group][attn_tp_rank][pp_rank] = { + "rank_ip": rank_ip, + "rank_port": rank_port, + } + logger.debug( + f"Register prefill bootstrap: DP{dp_group} TP{attn_tp_rank} PP{pp_rank} with rank_ip: {rank_ip} and rank_port: {rank_port}" + ) + + return web.Response(text="OK", status=200) + + async def _handle_route_get(self, request: web.Request): + engine_rank = request.query.get("engine_rank") + target_dp_group = request.query.get("target_dp_group") + target_pp_rank = request.query.get("target_pp_rank") + if not engine_rank or not target_dp_group or not target_pp_rank: + return web.Response(text="Missing inputs for bootstrap server.", status=400) + + # Currently we use engine_rank == -1 and target_dp_group == -1 to sync dp size + if ( + int(engine_rank) == -1 + and int(target_dp_group) == -1 + and int(target_pp_rank) == -1 + ): + prefill_parallel_info = { + "prefill_attn_tp_size": self.attn_tp_size, + "prefill_dp_size": self.dp_size, + "prefill_pp_size": self.pp_size, + } + return web.json_response(prefill_parallel_info, status=200) + + # Find corresponding prefill info + async with self.lock: + bootstrap_info = self.prefill_port_table[int(target_dp_group)][ + int(engine_rank) + ][int(target_pp_rank)] + + if bootstrap_info is not None: + return web.json_response(bootstrap_info, status=200) + else: + return web.Response(text="Bootstrap info not Found", status=404) + + def _run_server(self): + try: + # Event Loop + self._loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._loop) + + access_log = None + if logging.getLogger(__name__).getEffectiveLevel() <= logging.DEBUG: + access_log = self.app.logger + + self._runner = web.AppRunner(self.app, access_log=access_log) + self._loop.run_until_complete(self._runner.setup()) + + site = web.TCPSite(self._runner, host=self.host, port=self.port) + self._loop.run_until_complete(site.start()) + self._loop.run_forever() + except Exception as e: + logger.error(f"Server error: {str(e)}") + finally: + # Cleanup + self._loop.run_until_complete(self._runner.cleanup()) + self._loop.close() + + def close(self): + """Shutdown""" + if self._loop is not None and self._loop.is_running(): + self._loop.call_soon_threadsafe(self._loop.stop) + logger.info("Stopping server loop...") + + if self.thread.is_alive(): + self.thread.join(timeout=2) + logger.info("Server thread stopped") + + def poll(self) -> KVPoll: ... diff --git a/python/sglang/srt/disaggregation/mooncake/transfer_engine.py b/python/sglang/srt/disaggregation/mooncake/transfer_engine.py new file mode 100644 index 000000000..54657bb46 --- /dev/null +++ b/python/sglang/srt/disaggregation/mooncake/transfer_engine.py @@ -0,0 +1,164 @@ +import logging +from typing import List, Optional + +from sglang.srt.utils import get_bool_env_var, get_free_port, maybe_wrap_ipv6_address + +logger = logging.getLogger(__name__) + + +class MooncakeTransferEngine: + + def __init__(self, hostname: str, gpu_id: int, ib_device: Optional[str] = None): + try: + from mooncake.engine import TransferEngine + except ImportError as e: + raise ImportError( + "Please install mooncake by following the instructions at " + "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501 + "to run SGLang with MooncakeTransferEngine." + ) from e + + self.engine = TransferEngine() + self.hostname = hostname + self.gpu_id = gpu_id + self.ib_device = ib_device + + self.initialize( + hostname=self.hostname, + device_name=self.ib_device, + ) + self.session_id = ( + f"{maybe_wrap_ipv6_address(self.hostname)}:{self.engine.get_rpc_port()}" + ) + + def register(self, ptr, length): + try: + ret_value = self.engine.register_memory(ptr, length) + except Exception: + # Mark register as failed + ret_value = -1 + + if ret_value != 0: + logger.debug("Mooncake memory registration %s failed.", ptr) + + def deregister(self, ptr): + try: + ret_value = self.engine.unregister_memory(ptr) + except Exception: + # Mark deregister as failed + ret_value = -1 + + if ret_value != 0: + logger.debug("Mooncake memory deregistration %s failed.", ptr) + + def batch_register(self, ptrs: List[int], lengths: List[int]) -> int: + """Batch register multiple memory regions.""" + try: + ret_value = self.engine.batch_register_memory(ptrs, lengths) + except Exception: + # Mark batch register as failed + ret_value = -1 + if not hasattr(self.engine, "batch_register_memory"): + raise RuntimeError( + "Mooncake's batch register requires a newer version of mooncake-transfer-engine. " + "Please upgrade Mooncake." + ) + + if ret_value != 0: + logger.debug("Mooncake batch memory registration failed.") + return ret_value + + def batch_deregister(self, ptrs: List[int]) -> int: + """Batch deregister multiple memory regions.""" + try: + ret_value = self.engine.batch_unregister_memory(ptrs) + except Exception: + # Mark batch deregister as failed + ret_value = -1 + + if ret_value != 0: + logger.debug("Mooncake batch memory deregistration failed.") + return ret_value + + def initialize( + self, + hostname: str, + device_name: Optional[str], + ) -> None: + """Initialize the mooncake instance.""" + if get_bool_env_var("ENABLE_ASCEND_TRANSFER_WITH_MOONCAKE", "false"): + hostname += f":{get_free_port()}:npu_{self.gpu_id}" + ret_value = self.engine.initialize( + hostname, + "P2PHANDSHAKE", + "ascend", + device_name if device_name is not None else "", + ) + else: + ret_value = self.engine.initialize( + hostname, + "P2PHANDSHAKE", + "rdma", + device_name if device_name is not None else "", + ) + if ret_value != 0: + logger.error("Mooncake Transfer Engine initialization failed.") + raise RuntimeError("Mooncake Transfer Engine initialization failed.") + + def transfer_sync( + self, session_id: str, buffer: int, peer_buffer_address: int, length: int + ) -> int: + """Synchronously transfer data to the specified address.""" + try: + # the first time: based on session_id (which contains remote_ip) to construct a queue pair, and cache the queue pair + # later: based on the cached queue pair to send data + ret = self.engine.transfer_sync_write( + session_id, buffer, peer_buffer_address, length + ) + except Exception: + # Mark transfer request as failed + ret = -1 + + if ret < 0: + # Do not raise an exception here, since some transfer requests fail should be accepted and the execution thread should not be stopped. + logger.debug( + "Failed to transfer data from %s to %s - %s.", + buffer, + session_id, + peer_buffer_address, + ) + + return ret + + def batch_transfer_sync( + self, + session_id: str, + buffers: List[int], + peer_buffer_addresses: List[int], + lengths: List[int], + ) -> int: + """Synchronously transfer data to the specified addresses in batches.""" + try: + ret = self.engine.batch_transfer_sync_write( + session_id, buffers, peer_buffer_addresses, lengths + ) + except Exception: + ret = -1 + # Inform user to upgrade mooncake-transfer-engine >= 0.3.4.post2 + if not hasattr(self.engine, "batch_transfer_sync_write"): + raise RuntimeError( + "Mooncake's batch transfer requires mooncake-transfer-engine >= 0.3.4.post2. " + "Please upgrade Mooncake by 'pip install mooncake-transfer-engine --upgrade'" + ) + + if ret < 0: + logger.debug( + "Failed to batch transfer data. Buffers: %s, Session: %s, Peer addresses: %s", + buffers, + session_id, + peer_buffer_addresses, + ) + return ret + + def get_session_id(self): + return self.session_id diff --git a/python/sglang/srt/disaggregation/nixl/__init__.py b/python/sglang/srt/disaggregation/nixl/__init__.py new file mode 100644 index 000000000..4df7baba2 --- /dev/null +++ b/python/sglang/srt/disaggregation/nixl/__init__.py @@ -0,0 +1,6 @@ +from sglang.srt.disaggregation.nixl.conn import ( + NixlKVBootstrapServer, + NixlKVManager, + NixlKVReceiver, + NixlKVSender, +) diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py new file mode 100644 index 000000000..c911319ea --- /dev/null +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -0,0 +1,696 @@ +from __future__ import annotations + +import asyncio +import dataclasses +import logging +import queue +import socket +import struct +import threading +import uuid +from collections import defaultdict +from functools import cache +from typing import Dict, List, Optional, Set, Tuple, TypeAlias, Union + +import numpy as np +import numpy.typing as npt +import requests +import zmq +from aiohttp import web + +from sglang.srt.disaggregation.base.conn import BaseKVSender, KVArgs, KVPoll +from sglang.srt.disaggregation.common.conn import ( + CommonKVBootstrapServer, + CommonKVManager, + CommonKVReceiver, +) +from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import ( + format_tcp_address, + get_local_ip_auto, + is_valid_ipv6_address, +) + +logger = logging.getLogger(__name__) + +GUARD = "NixlMsgGuard".encode("ascii") + + +@dataclasses.dataclass +class TransferInfo: + """Contains indices for a transfer, sent by KVReceiver. Received by prefill bootstrap thread.""" + + room: int + endpoint: str + dst_port: int + agent_name: str + dst_kv_indices: npt.NDArray[np.int32] + dst_aux_index: int + required_dst_info_num: int + + def is_dummy(self): + return self.dst_kv_indices.size == 0 + + @classmethod + def from_zmq(cls, msg: List[bytes]): + return cls( + room=int(msg[0].decode("ascii")), + endpoint=msg[1].decode("ascii"), + dst_port=int(msg[2].decode("ascii")), + agent_name=msg[3].decode("ascii"), + dst_kv_indices=np.frombuffer(msg[4], dtype=np.int32), + dst_aux_index=int(msg[5].decode("ascii")), + required_dst_info_num=int(msg[6].decode("ascii")), + ) + + +@dataclasses.dataclass +class KVArgsRegisterInfo: + """Contains base pointers and other info which only needs to be sent once by KVReceiver. Received by prefill bootstrap thread.""" + + room: str + endpoint: str + dst_port: int + agent_name: str + agent_metadata: bytes + dst_kv_ptrs: list[int] + dst_aux_ptrs: list[int] + gpu_id: int + decode_tp_size: int + decode_tp_rank: int + dst_kv_item_len: int + + @classmethod + def from_zmq(cls, msg: List[bytes]): + return cls( + room=str(msg[0].decode("ascii")), + endpoint=msg[1].decode("ascii"), + dst_port=int(msg[2].decode("ascii")), + agent_name=msg[3].decode("ascii"), + agent_metadata=msg[4], + dst_kv_ptrs=list(struct.unpack(f"{len(msg[5])//8}Q", msg[5])), + dst_aux_ptrs=list(struct.unpack(f"{len(msg[6])//8}Q", msg[6])), + gpu_id=int(msg[7].decode("ascii")), + decode_tp_size=int(msg[8].decode("ascii")), + decode_tp_rank=int(msg[9].decode("ascii")), + dst_kv_item_len=int(msg[10].decode("ascii")), + ) + + +@dataclasses.dataclass +class TransferStatus: + """Used by KV Receiver to know when a transfer is done.""" + + # KV chunk IDs that have been received. + received_kvs: Set[int] = dataclasses.field(default_factory=set) + # Number of kv chunks to expect, will know this after last chunk is received. + num_kvs_expected: Optional[int] = None + # Whether aux data has been received. + received_aux: bool = False + + def is_done(self): + if self.num_kvs_expected is None: + return False + return self.num_kvs_expected == len(self.received_kvs) and self.received_aux + + +class NixlKVManager(CommonKVManager): + def __init__( + self, + args: KVArgs, + disaggregation_mode: DisaggregationMode, + server_args: ServerArgs, + is_mla_backend: Optional[bool] = False, + ): + super().__init__(args, disaggregation_mode, server_args, is_mla_backend) + try: + from nixl._api import nixl_agent + except ImportError as e: + raise ImportError( + "Please install NIXL by following the instructions at " + "https://github.com/ai-dynamo/nixl/blob/main/README.md " + "to run SGLang with NixlTransferEngine." + ) from e + self.agent = nixl_agent(str(uuid.uuid4())) + self.local_ip = get_local_ip_auto() + self.server_socket = zmq.Context().socket(zmq.PULL) + if is_valid_ipv6_address(self.local_ip): + self.server_socket.setsockopt(zmq.IPV6, 1) + self.register_buffer_to_engine() + + if self.disaggregation_mode == DisaggregationMode.PREFILL: + self.request_status: Dict[int, KVPoll] = {} + self.transfer_infos: Dict[int, Dict[str, TransferInfo]] = {} + self.decode_kv_args_table: Dict[str, KVArgsRegisterInfo] = {} + self._start_bootstrap_thread() + elif self.disaggregation_mode == DisaggregationMode.DECODE: + self.transfer_statuses: Dict[int, TransferStatus] = defaultdict( + TransferStatus + ) + else: + raise ValueError( + f"Unsupported DisaggregationMode: {self.disaggregation_mode}" + ) + + def check_status(self, bootstrap_room: int): + return self.request_status[bootstrap_room] + + def update_status(self, bootstrap_room: int, status: KVPoll): + if bootstrap_room not in self.request_status: + self.request_status[bootstrap_room] = status + else: + # NOTE: The prefill engine could recv bootstrapping first + self.request_status[bootstrap_room] = max( + self.request_status[bootstrap_room], status + ) + + def register_buffer_to_engine(self): + kv_addrs = [] + for kv_data_ptr, kv_data_len in zip( + self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens + ): + kv_addrs.append((kv_data_ptr, kv_data_len, self.kv_args.gpu_id, "")) + self.kv_descs = self.agent.register_memory(kv_addrs, "VRAM") + logger.debug(f"Register kv tensors, len(kv_addr)= {len(kv_addrs)}") + if not self.kv_descs: + raise Exception("NIXL memory registration failed for kv tensors") + aux_addrs = [] + for aux_data_ptr, aux_data_len in zip( + self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens + ): + aux_addrs.append((aux_data_ptr, aux_data_len, 0, "")) + self.aux_descs = self.agent.register_memory(aux_addrs, "DRAM") + logger.debug(f"Register aux tensors, len(aux_addrs)= {len(aux_addrs)}") + if not self.aux_descs: + raise Exception("NIXL memory registration failed for aux tensors") + + def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo): + agent_name = decode_kv_args.agent_name + if agent_name in self.decode_kv_args_table: + logger.info(f"Peer {agent_name} was already registered, ignoring.") + return + self.decode_kv_args_table[agent_name] = decode_kv_args + self.agent.add_remote_agent(decode_kv_args.agent_metadata) + + def send_kvcache( + self, + peer_name: str, + prefill_kv_indices: npt.NDArray[np.int32], + dst_kv_ptrs: list[int], + dst_kv_indices: npt.NDArray[np.int32], + dst_gpu_id: int, + notif: str, + ): + # group by indices + prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( + prefill_kv_indices, dst_kv_indices + ) + + logger.debug(f"sending kvcache to {peer_name} with notif {notif}") + # Make descs + num_layers = len(self.kv_args.kv_data_ptrs) + src_addrs = [] + dst_addrs = [] + for layer_id in range(num_layers): + src_ptr = self.kv_args.kv_data_ptrs[layer_id] + dst_ptr = dst_kv_ptrs[layer_id] + item_len = self.kv_args.kv_item_lens[layer_id] + + for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks): + src_addr = src_ptr + int(prefill_index[0]) * item_len + dst_addr = dst_ptr + int(decode_index[0]) * item_len + length = item_len * len(prefill_index) + src_addrs.append((src_addr, length, self.kv_args.gpu_id)) + dst_addrs.append((dst_addr, length, dst_gpu_id)) + + logger.debug( + f"len(src_addrs): before group: {len(prefill_kv_indices)}, after group: {len(src_addrs)}" + ) + src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM") + dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM") + # Transfer data + xfer_handle = self.agent.initialize_xfer( + "WRITE", + src_descs, + dst_descs, + peer_name, + notif.encode("ascii"), # type: ignore + ) + if not xfer_handle: + raise Exception("KVSender failed to create transfer") + state = self.agent.transfer(xfer_handle) + if state == "ERR": + raise Exception("KVSender failed to post transfer") + return xfer_handle + + def send_kvcache_slice( + self, + peer_name: str, + prefill_kv_indices: npt.NDArray[np.int32], + dst_kv_ptrs: list[int], + dst_kv_indices: npt.NDArray[np.int32], + dst_gpu_id: int, + notif: str, + prefill_tp_size: int, + decode_tp_size: int, + decode_tp_rank: int, + dst_kv_item_len: int, + ): + # Get configuration from kv_args + local_tp_rank_in_group = self.kv_args.engine_rank % prefill_tp_size + dst_tp_rank_in_group = decode_tp_rank % decode_tp_size + num_kv_heads = self.kv_args.kv_head_num + + # Calculate head distribution + src_heads_per_rank = num_kv_heads + dst_heads_per_rank = num_kv_heads * prefill_tp_size // decode_tp_size + + src_kv_item_len = self.kv_args.kv_item_lens[0] + page_size = self.kv_args.page_size + + bytes_per_head_slice_to_send = ( + dst_kv_item_len // page_size // dst_heads_per_rank + ) + + # Determine which heads to send + if prefill_tp_size > decode_tp_size: + # Multiple prefill ranks to one decode rank + src_head_start_offset = 0 + num_heads_to_send = src_heads_per_rank + dst_head_start_offset = local_tp_rank_in_group * src_heads_per_rank + else: + # Send KVCache from 1 prefill instance to multiple decode instances + src_head_start_offset = ( + dst_tp_rank_in_group * dst_heads_per_rank + ) % src_heads_per_rank + num_heads_to_send = dst_heads_per_rank + dst_head_start_offset = 0 + + # Create transfer descriptors + src_addrs = [] + dst_addrs = [] + + bytes_per_token_on_prefill = src_kv_item_len // page_size + bytes_per_token_on_decode = dst_kv_item_len // page_size + + num_kv_layers = len(self.kv_args.kv_data_ptrs) // 2 + src_k_ptrs = self.kv_args.kv_data_ptrs[:num_kv_layers] + src_v_ptrs = self.kv_args.kv_data_ptrs[num_kv_layers:] + dst_k_ptrs = dst_kv_ptrs[0 : len(src_k_ptrs)] + dst_v_ptrs = dst_kv_ptrs[num_kv_layers : num_kv_layers + len(src_v_ptrs)] + + # Calculate precise byte offset and length for the sub-slice within the token + src_head_slice_offset = src_head_start_offset * bytes_per_head_slice_to_send + dst_head_slice_offset = dst_head_start_offset * bytes_per_head_slice_to_send + heads_bytes_per_token_to_send = num_heads_to_send * bytes_per_head_slice_to_send + + src_dst_ptr_pairs = [ + ( + src_k_ptrs[layer_id], + dst_k_ptrs[layer_id], + ) + for layer_id in range(len(src_k_ptrs)) + ] + [ + ( + src_v_ptrs[layer_id], + dst_v_ptrs[layer_id], + ) + for layer_id in range(len(src_v_ptrs)) + ] + + src_addrs = [] + dst_addrs = [] + + # Calculate strides for a single token slot + bytes_per_token_on_prefill = src_kv_item_len // page_size + bytes_per_token_on_decode = dst_kv_item_len // page_size + + for src_ptr, dst_ptr in src_dst_ptr_pairs: + for i in range(len(prefill_kv_indices)): + prefill_page_idx = int(prefill_kv_indices[i]) + decode_page_idx = int(dst_kv_indices[i]) + + # Get the starting addresses for the current src and dst pages + src_page_start_addr = src_ptr + prefill_page_idx * src_kv_item_len + dst_page_start_addr = dst_ptr + decode_page_idx * dst_kv_item_len + + # Iterate through each valid token slot within the current page + for token_slot_in_page in range(page_size): + # Calculate the start address of the current token slot + src_token_slot_start_addr = ( + src_page_start_addr + + token_slot_in_page * bytes_per_token_on_prefill + ) + dst_token_slot_start_addr = ( + dst_page_start_addr + + token_slot_in_page * bytes_per_token_on_decode + ) + + # Calculate final src and dst addresses by applying head-slice offsets + src_slice_addr = src_token_slot_start_addr + src_head_slice_offset + dst_slice_addr = dst_token_slot_start_addr + dst_head_slice_offset + + src_addrs.append( + ( + src_slice_addr, + heads_bytes_per_token_to_send, + self.kv_args.gpu_id, + ) + ) + dst_addrs.append( + (dst_slice_addr, heads_bytes_per_token_to_send, dst_gpu_id) + ) + + # Use NIXL agent for transfer + src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM") + dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM") + + xfer_handle = self.agent.initialize_xfer( + "WRITE", src_descs, dst_descs, peer_name, notif.encode("ascii") + ) + if not xfer_handle: + raise Exception("Failed to create sliced KV transfer") + + state = self.agent.transfer(xfer_handle) + if state == "ERR": + raise Exception("Failed to post sliced KV transfer") + + return xfer_handle + + def send_aux( + self, + peer_name: str, + prefill_aux_index: int, + dst_aux_ptrs: list[int], + dst_aux_index: int, + notif: str, + ): + # Make descs + aux_item_len = self.kv_args.aux_item_lens[0] + prefill_aux_addr = ( + self.kv_args.aux_data_ptrs[0] + prefill_aux_index * aux_item_len + ) + decode_aux_addr = dst_aux_ptrs[0] + dst_aux_index * aux_item_len + src_addrs = [(prefill_aux_addr, aux_item_len, 0)] + dst_addrs = [(decode_aux_addr, aux_item_len, 0)] + src_descs = self.agent.get_xfer_descs(src_addrs, "DRAM") + dst_descs = self.agent.get_xfer_descs(dst_addrs, "DRAM") + # Transfer data + xfer_handle = self.agent.initialize_xfer( + "WRITE", + src_descs, + dst_descs, + peer_name, + notif.encode("ascii"), # type: ignore + ) + if not xfer_handle: + raise Exception("KVSender failed to create transfer") + state = self.agent.transfer(xfer_handle) + if state == "ERR": + raise Exception("KVSender failed to post transfer") + return xfer_handle + + def add_transfer_request( + self, + bootstrap_room: int, + kv_indices: npt.NDArray[np.int32], + index_slice: slice, + is_last: bool, + chunk_id: int, + aux_index: Optional[int] = None, + ): + assert self.disaggregation_mode == DisaggregationMode.PREFILL + assert not is_last or (is_last and aux_index is not None) + + reqs_to_be_processed = self.transfer_infos[bootstrap_room].values() + handles = [] + for req in reqs_to_be_processed: + assert bootstrap_room == req.room + if req.is_dummy(): + continue + + chunked_dst_kv_indice = req.dst_kv_indices[index_slice] + assert len(chunked_dst_kv_indice) == len(kv_indices) + assert req.agent_name in self.decode_kv_args_table + + notif = "_".join([str(req.room), "kv", str(chunk_id), str(int(is_last))]) + decode_tp_size = self.decode_kv_args_table[req.agent_name].decode_tp_size + + if decode_tp_size == self.tp_size: + kv_xfer_handle = self.send_kvcache( + req.agent_name, + kv_indices, + self.decode_kv_args_table[req.agent_name].dst_kv_ptrs, + chunked_dst_kv_indice, + self.decode_kv_args_table[req.agent_name].gpu_id, + notif, + ) + else: + kv_xfer_handle = self.send_kvcache_slice( + req.agent_name, + kv_indices, + self.decode_kv_args_table[req.agent_name].dst_kv_ptrs, + chunked_dst_kv_indice, + self.decode_kv_args_table[req.agent_name].gpu_id, + notif, + prefill_tp_size=self.tp_size, + decode_tp_size=decode_tp_size, + decode_tp_rank=self.decode_kv_args_table[ + req.agent_name + ].decode_tp_rank, + dst_kv_item_len=self.decode_kv_args_table[ + req.agent_name + ].dst_kv_item_len, + ) + + handles.append(kv_xfer_handle) + # Only the last chunk we need to send the aux data. + if is_last: + assert aux_index is not None + aux_xfer_handle = self.send_aux( + req.agent_name, + aux_index, + self.decode_kv_args_table[req.agent_name].dst_aux_ptrs, + req.dst_aux_index, + str(req.room) + "_aux", + ) + handles.append(aux_xfer_handle) + if is_last: + del self.transfer_infos[bootstrap_room] + return handles + + def update_transfer_status(self): + # Process notifications from received transfers. + notif_map = self.agent.get_new_notifs() + for peer_name, messages in notif_map.items(): + # We could also check that self.bootstrap_info['agent_name'] matches + # the message sender. But the bootstrap room alone should be + # sufficient to map the status. + for msg in messages: + components = msg.decode("ascii").split("_") + room = int(components[0]) + if components[1] == "kv": + chunk_id = int(components[2]) + is_last = bool(int(components[3])) + self.transfer_statuses[room].received_kvs.add(chunk_id) + if is_last: + self.transfer_statuses[room].num_kvs_expected = chunk_id + 1 + elif components[1] == "aux": + self.transfer_statuses[room].received_aux = True + + def check_transfer_done(self, room: int): + if room not in self.transfer_statuses: + return False + return self.transfer_statuses[room].is_done() + + def _bind_server_socket(self): + self.server_socket.bind(format_tcp_address(self.local_ip, self.rank_port)) + + def _start_bootstrap_thread(self): + self._bind_server_socket() + + def bootstrap_thread(): + """This thread recvs transfer info from the decode engine""" + while True: + waiting_req_bytes = self.server_socket.recv_multipart() + logger.debug( + f"Received multipart with total byte size {sum(len(x) for x in waiting_req_bytes)}" + ) + assert ( + waiting_req_bytes[0] == GUARD + ), f"First message should be {GUARD}. Foreign traffic?" + waiting_req_bytes = waiting_req_bytes[1:] + room = waiting_req_bytes[0].decode("ascii") + agent_name = waiting_req_bytes[3].decode("ascii") + if room == "None": + # Register new peer and save KV base pointers. + self._add_remote_peer( + KVArgsRegisterInfo.from_zmq(waiting_req_bytes) + ) + logger.debug(f"Register KVArgs from {agent_name} successfully") + continue + room = int(room) + if room not in self.transfer_infos: + self.transfer_infos[room] = {} + self.transfer_infos[room][agent_name] = TransferInfo.from_zmq( + waiting_req_bytes + ) + required_dst_info_num = self.transfer_infos[room][ + agent_name + ].required_dst_info_num + logger.debug(f"got info {room=} {agent_name=} {required_dst_info_num=}") + if len(self.transfer_infos[room]) == required_dst_info_num: + logger.debug(f"{room=} is bootstrapped") + self.update_status(room, KVPoll.WaitingForInput) + + threading.Thread(target=bootstrap_thread).start() + + +class NixlKVSender(BaseKVSender): + + def __init__( + self, + mgr: NixlKVManager, + bootstrap_addr: str, + bootstrap_room: int, + dest_tp_ranks: List[int], + pp_rank: int, + ): + self.kv_mgr = mgr + self.bootstrap_room = bootstrap_room + self.aux_index = None + self.bootstrap_server_url = bootstrap_addr + self.xfer_handles = [] + self.has_sent = False + self.chunk_id = 0 + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) + # inner state + self.curr_idx = 0 + + def init(self, num_kv_indices: int, aux_index: Optional[int] = None): + self.num_kv_indices = num_kv_indices + self.aux_index = aux_index + + def send( + self, + kv_indices: npt.NDArray[np.int32], + ): + index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices)) + self.curr_idx += len(kv_indices) + is_last = self.curr_idx == self.num_kv_indices + + new_xfer_handles = self.kv_mgr.add_transfer_request( + self.bootstrap_room, + kv_indices, + index_slice, + is_last, + self.chunk_id, + self.aux_index, + ) + self.xfer_handles.extend(new_xfer_handles) + self.chunk_id += 1 + if is_last: + self.has_sent = True + del self.kv_mgr.request_status[self.bootstrap_room] + + def poll(self) -> KVPoll: + if not self.has_sent: + return self.kv_mgr.check_status(self.bootstrap_room) + states = [self.kv_mgr.agent.check_xfer_state(x) for x in self.xfer_handles] + if all([x == "DONE" for x in states]): + return KVPoll.Success # type: ignore + if any([x == "ERR" for x in states]): + raise Exception("KVSender transfer encountered an error.") + return KVPoll.WaitingForInput # type: ignore + + def failure_exception(self): + raise Exception("Fake KVSender Exception") + + +class NixlKVReceiver(CommonKVReceiver): + def __init__( + self, + mgr: NixlKVManager, + bootstrap_addr: str, + bootstrap_room: Optional[int] = None, + prefill_dp_rank: Optional[int] = None, + ): + self.started_transfer = False + self.conclude_state = None + super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank) + + def init(self, kv_indices: npt.NDArray[np.int32], aux_index: Optional[int] = None): + for bootstrap_info in self.bootstrap_infos: + logger.debug( + f"Fetched bootstrap info: {bootstrap_info} for engine rank: {self.kv_mgr.kv_args.engine_rank}" + ) + sock, lock = self._connect_to_bootstrap_server(bootstrap_info) + is_dummy = bootstrap_info["is_dummy"] + logger.debug( + f"Sending to prefill server with bootstrap room {self.bootstrap_room} {is_dummy=}" + ) + with lock: + sock.send_multipart( + [ + GUARD, + str(self.bootstrap_room).encode("ascii"), + self.kv_mgr.local_ip.encode("ascii"), + str(self.kv_mgr.rank_port).encode("ascii"), + self.kv_mgr.agent.name.encode("ascii"), + kv_indices.tobytes() if not is_dummy else b"", + str(aux_index).encode("ascii"), + str(self.required_dst_info_num).encode("ascii"), + ] + ) + + self.started_transfer = True + + def poll(self) -> KVPoll: + if self.conclude_state is not None: + return self.conclude_state + if not self.started_transfer: + return KVPoll.WaitingForInput # type: ignore + + self.kv_mgr.update_transfer_status() + if self.kv_mgr.check_transfer_done(self.bootstrap_room): # type: ignore + self.conclude_state = KVPoll.Success + del self.kv_mgr.transfer_statuses[self.bootstrap_room] + return KVPoll.Success # type: ignore + return KVPoll.WaitingForInput # type: ignore + + def _register_kv_args(self): + for bootstrap_info in self.bootstrap_infos: + sock, lock = self._connect_to_bootstrap_server(bootstrap_info) + packed_kv_data_ptrs = b"".join( + struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs + ) + packed_aux_data_ptrs = b"".join( + struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs + ) + + with lock: + sock.send_multipart( + [ + GUARD, + "None".encode("ascii"), + self.kv_mgr.local_ip.encode("ascii"), + str(self.kv_mgr.rank_port).encode("ascii"), + self.kv_mgr.agent.name.encode("ascii"), + self.kv_mgr.agent.get_agent_metadata(), + packed_kv_data_ptrs, + packed_aux_data_ptrs, + str(self.kv_mgr.kv_args.gpu_id).encode("ascii"), + str(self.kv_mgr.kv_args.decode_tp_size).encode("ascii"), + str(self.kv_mgr.kv_args.engine_rank).encode("ascii"), + str(self.kv_mgr.kv_args.kv_item_lens[0]).encode("ascii"), + ] + ) + + def failure_exception(self): + raise Exception("Fake KVReceiver Exception") + + +class NixlKVBootstrapServer(CommonKVBootstrapServer): + pass diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py new file mode 100644 index 000000000..b70748250 --- /dev/null +++ b/python/sglang/srt/disaggregation/prefill.py @@ -0,0 +1,867 @@ +""" +Life cycle of a request in the prefill server + +1. Bootstrap Queue + a. Initialize a sender for each request + b. Use the queue to store requests whose bootstrap (handshake and preallocation) has not finished + c. Poll senders to check bootstrap state + d. Once bootstrap is complete, move request to Waiting Queue + +2. Waiting Queue + a. Use PrefillAdder to pop requests + b. Run forward + c. Add the request to Inflight Queue + +3. Inflight Queue + a. Poll (non-blocking) the sender of the request + b. Once the transfer has finished, return the request +""" + +from __future__ import annotations + +import logging +import threading +from collections import deque +from http import HTTPStatus +from typing import TYPE_CHECKING, List, Optional, Type + +import torch + +from sglang.srt.disaggregation.base import BaseKVManager, KVPoll +from sglang.srt.disaggregation.utils import ( + FAKE_BOOTSTRAP_HOST, + DisaggregationMode, + KVClassType, + MetadataBuffers, + ReqToMetadataIdxAllocator, + TransferBackend, + get_kv_class, + is_mla_backend, + kv_to_page_indices, + kv_to_page_num, + poll_and_all_reduce, + prepare_abort, +) +from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch +from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors +from sglang.srt.utils import ( + DynamicGradMode, + broadcast_pyobj, + point_to_point_pyobj, + require_mlp_sync, +) + +if TYPE_CHECKING: + from torch.distributed import ProcessGroup + + from sglang.srt.managers.scheduler import GenerationBatchResult, Scheduler + from sglang.srt.mem_cache.memory_pool import KVCache + +logger = logging.getLogger(__name__) + + +class PrefillBootstrapQueue: + """ + Store the requests in bootstrapping + """ + + def __init__( + self, + token_to_kv_pool: KVCache, + draft_token_to_kv_pool: Optional[KVCache], + req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator, + metadata_buffers: MetadataBuffers, + tp_rank: int, + tp_size: int, + gpu_id: int, + bootstrap_port: int, + gloo_group: ProcessGroup, + max_total_num_tokens: int, + decode_tp_size: int, + decode_dp_size: int, + scheduler: Scheduler, + pp_rank: int, + pp_size: int, + transfer_backend: TransferBackend, + ): + self.token_to_kv_pool = token_to_kv_pool + self.draft_token_to_kv_pool = draft_token_to_kv_pool + self.is_mla_backend = is_mla_backend(token_to_kv_pool) + self.metadata_buffers = metadata_buffers + self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator + self.tp_rank = tp_rank + self.tp_size = tp_size + self.decode_tp_size = decode_tp_size + self.decode_dp_size = decode_dp_size + self.pp_rank = pp_rank + self.pp_size = pp_size + self.gpu_id = gpu_id + self.bootstrap_port = bootstrap_port + self.queue: List[Req] = [] + self.gloo_group = gloo_group + self.max_total_num_tokens = max_total_num_tokens + self.scheduler = scheduler + self.transfer_backend = transfer_backend + self.kv_manager = self._init_kv_manager() + + def _init_kv_manager(self) -> BaseKVManager: + kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS) + kv_args = kv_args_class() + kv_args.engine_rank = self.tp_rank + kv_args.pp_rank = self.pp_rank + kv_args.system_dp_rank = self.scheduler.dp_rank + kv_args.decode_tp_size = self.decode_tp_size // self.decode_dp_size + kv_args.prefill_pp_size = self.pp_size + kv_args.prefill_start_layer = self.token_to_kv_pool.start_layer + kv_data_ptrs, kv_data_lens, kv_item_lens = ( + self.token_to_kv_pool.get_contiguous_buf_infos() + ) + + if self.draft_token_to_kv_pool is not None: + # We should also transfer draft model kv cache. The indices are + # always shared with a target model. + draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( + self.draft_token_to_kv_pool.get_contiguous_buf_infos() + ) + kv_data_ptrs += draft_kv_data_ptrs + kv_data_lens += draft_kv_data_lens + kv_item_lens += draft_kv_item_lens + + kv_args.kv_data_ptrs = kv_data_ptrs + kv_args.kv_data_lens = kv_data_lens + kv_args.kv_item_lens = kv_item_lens + if not self.is_mla_backend: + kv_args.kv_head_num = self.token_to_kv_pool.head_num + kv_args.page_size = self.token_to_kv_pool.page_size + + kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = ( + self.metadata_buffers.get_buf_infos() + ) + kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device + kv_args.gpu_id = self.scheduler.gpu_id + + kv_manager_class: Type[BaseKVManager] = get_kv_class( + self.transfer_backend, KVClassType.MANAGER + ) + kv_manager: BaseKVManager = kv_manager_class( + kv_args, + DisaggregationMode.PREFILL, + self.scheduler.server_args, + self.is_mla_backend, + ) + return kv_manager + + def add(self, req: Req, num_kv_heads: int) -> None: + if self._check_if_req_exceed_kv_capacity(req): + return + + if req.bootstrap_host == FAKE_BOOTSTRAP_HOST: + kv_sender_class = get_kv_class(TransferBackend.FAKE, KVClassType.SENDER) + else: + kv_sender_class = get_kv_class(self.transfer_backend, KVClassType.SENDER) + + dest_tp_ranks = [self.tp_rank] + + req.disagg_kv_sender = kv_sender_class( + mgr=self.kv_manager, + bootstrap_addr=f"{req.bootstrap_host}:{self.bootstrap_port}", + bootstrap_room=req.bootstrap_room, + dest_tp_ranks=dest_tp_ranks, + pp_rank=self.pp_rank, + ) + self._process_req(req) + self.queue.append(req) + + def extend(self, reqs: List[Req], num_kv_heads: int) -> None: + for req in reqs: + self.add(req, num_kv_heads) + + def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: + if len(req.origin_input_ids) > self.max_total_num_tokens: + message = f"Request {req.rid} exceeds the maximum number of tokens: {len(req.origin_input_ids)} > {self.max_total_num_tokens}" + logger.error(message) + prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) + self.scheduler.stream_output([req], req.return_logprob) + return True + return False + + def _process_req(self, req: Req) -> None: + """ + Set max_new_tokens = 1, so PrefillAdder memory estimation is accurate + """ + req.sampling_params.max_new_tokens = 1 + + def pop_bootstrapped( + self, + return_failed_reqs: bool = False, + rids_to_check: Optional[List[str]] = None, + ) -> List[Req]: + """ + pop the reqs which has finished bootstrapping + + return_failed_reqs: For PP, on rank 0, also return the failed reqs to notify the next rank + rids_to_check: For PP, on rank > 0, check the rids from the previous rank has consensus with the current rank. + """ + + bootstrapped_reqs = [] + failed_reqs = [] + indices_to_remove = set() + + if len(self.queue) == 0: + if return_failed_reqs is False: + return [] + else: + return [], [] + + polls = poll_and_all_reduce( + [req.disagg_kv_sender for req in self.queue], self.gloo_group + ) + + for i, (req, poll) in enumerate(zip(self.queue, polls)): + if rids_to_check is not None: + # if req not in reqs_info_to_check, skip + if req.rid not in rids_to_check: + continue + # Either waiting for input or failed + assert poll == KVPoll.WaitingForInput or poll == KVPoll.Failed + + if poll == KVPoll.Bootstrapping: + continue + elif poll == KVPoll.Failed: + error_message = f"Prefill bootstrap failed for request rank={self.tp_rank} {req.rid=} {req.bootstrap_room=}" + try: + req.disagg_kv_sender.failure_exception() + except Exception as e: + error_message += f" with exception {e}" + logger.error(error_message) + prepare_abort( + req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR + ) + self.scheduler.stream_output([req], req.return_logprob) + indices_to_remove.add(i) + failed_reqs.append(req) + if self.scheduler.enable_metrics: + self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() + continue + + # KV.WaitingForInput - init here + num_kv_indices = len(req.origin_input_ids) + if self.req_to_metadata_buffer_idx_allocator.available_size() == 0: + break + + req.metadata_buffer_index = ( + self.req_to_metadata_buffer_idx_allocator.alloc() + ) + assert req.metadata_buffer_index is not None + + num_pages = kv_to_page_num(num_kv_indices, self.token_to_kv_pool.page_size) + req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index) + bootstrapped_reqs.append(req) + indices_to_remove.add(i) + + self.queue = [ + entry for i, entry in enumerate(self.queue) if i not in indices_to_remove + ] + + if return_failed_reqs is False: + return bootstrapped_reqs + else: + return bootstrapped_reqs, failed_reqs + + +class SchedulerDisaggregationPrefillMixin: + """ + Mixin for Scheduler to handle disaggregation prefill + """ + + @torch.no_grad() + def event_loop_normal_disagg_prefill(self: Scheduler) -> None: + """A normal scheduler loop for prefill worker in disaggregation mode.""" + + while True: + recv_reqs = self.recv_requests() + self.process_input_requests(recv_reqs) + self.waiting_queue.extend( + self.disagg_prefill_bootstrap_queue.pop_bootstrapped() + ) + self.process_prefill_chunk() + batch = self.get_new_batch_prefill() + + if require_mlp_sync(self.server_args): + batch = self.prepare_mlp_sync_batch(batch) + self.cur_batch = batch + + if batch: + result = self.run_batch(batch) + self.process_batch_result_disagg_prefill(batch, result) + + if len(self.disagg_prefill_inflight_queue) > 0: + self.process_disagg_prefill_inflight_queue() + + if batch is None and len(self.disagg_prefill_inflight_queue) == 0: + self.self_check_during_idle() + + self.last_batch = batch + # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it + # Otherwise, it hangs under high concurrency + self.running_batch.batch_is_full = False + + @torch.no_grad() + def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: + self.result_queue = deque() + + while True: + recv_reqs = self.recv_requests() + self.process_input_requests(recv_reqs) + self.waiting_queue.extend( + self.disagg_prefill_bootstrap_queue.pop_bootstrapped() + ) + self.process_prefill_chunk() + batch = self.get_new_batch_prefill() + + if require_mlp_sync(self.server_args): + batch = self.prepare_mlp_sync_batch(batch) + self.cur_batch = batch + if batch: + result = self.run_batch(batch) + self.result_queue.append((batch.copy(), result)) + + if self.last_batch is None: + # Create a dummy first batch to start the pipeline for overlap schedule. + # It is now used for triggering the sampling_info_done event. + tmp_batch = ScheduleBatch( + reqs=None, + forward_mode=ForwardMode.DUMMY_FIRST, + next_batch_sampling_info=self.tp_worker.cur_sampling_info, + ) + self.set_next_batch_sampling_info_done(tmp_batch) + + if self.last_batch: + tmp_batch, tmp_result = self.result_queue.popleft() + tmp_batch.next_batch_sampling_info = ( + self.tp_worker.cur_sampling_info if batch else None + ) + self.process_batch_result_disagg_prefill(tmp_batch, tmp_result) + + if len(self.disagg_prefill_inflight_queue) > 0: + self.process_disagg_prefill_inflight_queue() + + if batch is None and len(self.disagg_prefill_inflight_queue) == 0: + self.self_check_during_idle() + + self.last_batch = batch + # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it + # Otherwise, it hangs under high concurrency + self.running_batch.batch_is_full = False + + def process_batch_result_disagg_prefill( + self: Scheduler, + batch: ScheduleBatch, + result: GenerationBatchResult, + launch_done: Optional[threading.Event] = None, + ) -> None: + """ + Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue + Adapted from process_batch_result_prefill + """ + ( + logits_output, + next_token_ids, + extend_input_len_per_req, + extend_logprob_start_len_per_req, + ) = ( + result.logits_output, + result.next_token_ids, + result.extend_input_len_per_req, + result.extend_logprob_start_len_per_req, + ) + + logprob_pt = 0 + # Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue + if self.enable_overlap: + # wait + logits_output, next_token_ids, _ = self.tp_worker.resolve_last_batch_result( + launch_done + ) + else: + next_token_ids = result.next_token_ids.tolist() + if batch.return_logprob: + if logits_output.next_token_logprobs is not None: + logits_output.next_token_logprobs = ( + logits_output.next_token_logprobs.tolist() + ) + if logits_output.input_token_logprobs is not None: + logits_output.input_token_logprobs = tuple( + logits_output.input_token_logprobs.tolist() + ) + + hidden_state_offset = 0 + for i, (req, next_token_id) in enumerate( + zip(batch.reqs, next_token_ids, strict=True) + ): + req: Req + if req.is_chunked <= 0: + # There is no output_ids for prefill + req.output_ids.append(next_token_id) + self.tree_cache.cache_unfinished_req(req) # update the tree and lock + self.disagg_prefill_inflight_queue.append(req) + if ( + logits_output is not None + and logits_output.hidden_states is not None + ): + last_hidden_index = ( + hidden_state_offset + extend_input_len_per_req[i] - 1 + ) + req.hidden_states_tensor = ( + logits_output.hidden_states[last_hidden_index].cpu().clone() + ) + hidden_state_offset += extend_input_len_per_req[i] + else: + req.hidden_states_tensor = None + if req.return_logprob: + assert extend_logprob_start_len_per_req is not None + assert extend_input_len_per_req is not None + extend_logprob_start_len = extend_logprob_start_len_per_req[i] + extend_input_len = extend_input_len_per_req[i] + num_input_logprobs = extend_input_len - extend_logprob_start_len + self.add_logprob_return_values( + i, + req, + logprob_pt, + next_token_ids, + num_input_logprobs, + logits_output, + ) + logprob_pt += num_input_logprobs + self.send_kv_chunk(req, last_chunk=True) + + if req.grammar is not None: + # FIXME: this try-except block is for handling unexpected xgrammar issue. + try: + req.grammar.accept_token(next_token_id) + except ValueError as e: + # Grammar accept_token can raise ValueError if the token is not in the grammar. + # This can happen if the grammar is not set correctly or the token is invalid. + error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}" + self.tree_cache.cache_finished_req(req) + prepare_abort( + req, + error_message, + status_code=HTTPStatus.INTERNAL_SERVER_ERROR, + ) + req.grammar.finished = req.finished() + else: + # being chunked reqs' prefill is not finished + req.is_chunked -= 1 + + if req.return_logprob: + extend_logprob_start_len = extend_logprob_start_len_per_req[i] + extend_input_len = extend_input_len_per_req[i] + if extend_logprob_start_len < extend_input_len: + # Update input logprobs. + num_input_logprobs = extend_input_len - extend_logprob_start_len + self.add_input_logprob_return_values( + i, + req, + logits_output, + logprob_pt, + num_input_logprobs, + last_prefill_chunk=False, + ) + logprob_pt += num_input_logprobs + + if self.enable_overlap: + self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx) + + # We need to remove the sync in the following function for overlap schedule. + self.set_next_batch_sampling_info_done(batch) + self.maybe_send_health_check_signal() + + def process_disagg_prefill_inflight_queue( + self: Scheduler, rids_to_check: Optional[List[str]] = None + ) -> List[Req]: + """ + Poll the requests in the middle of transfer. If done, return the request. + rids_to_check: For PP, on rank > 0, check the rids from the previous rank has consensus with the current rank. + """ + if len(self.disagg_prefill_inflight_queue) == 0: + return [] + + done_reqs = [] + + polls = poll_and_all_reduce( + [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue], + self.attn_tp_cpu_group, + ) + + undone_reqs: List[Req] = [] + # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue + for req, poll in zip(self.disagg_prefill_inflight_queue, polls): + + if rids_to_check is not None: + if req.rid not in rids_to_check: + undone_reqs.append(req) + continue + + assert poll == KVPoll.Success or poll == KVPoll.Failed + + if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: + undone_reqs.append(req) + elif poll == KVPoll.Success: # transfer done + self.tree_cache.cache_finished_req(req) # unlock the tree + req.finished_reason = FINISH_LENGTH(length=0) + # FIXME: clean up req's data in transfer engine + if hasattr(req.disagg_kv_sender, "clear"): + req.disagg_kv_sender.clear() + done_reqs.append(req) + elif poll == KVPoll.Failed: + error_message = f"Prefill transfer failed for request rank={self.tp_rank} {req.rid=} {req.bootstrap_room=}" + try: + req.disagg_kv_sender.failure_exception() + except Exception as e: + error_message += f" with exception {e}" + logger.warning(error_message) + self.tree_cache.cache_finished_req(req) # unlock the tree + prepare_abort( + req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR + ) + done_reqs.append(req) + if self.enable_metrics: + self.metrics_collector.increment_transfer_failed_reqs() + else: + assert False, f"Unexpected polling state {poll=}" + + # Stream requests which have finished transfer + self.stream_output( + done_reqs, + any(req.return_logprob for req in done_reqs), + None, + ) + for req in done_reqs: + req: Req + self.req_to_metadata_buffer_idx_allocator.free(req.metadata_buffer_index) + req.metadata_buffer_index = -1 + + self.disagg_prefill_inflight_queue = undone_reqs + + return done_reqs + + def get_transferred_rids(self: Scheduler) -> List[str]: + """ + Used by PP, get the transferred rids but **do not pop** + """ + polls = poll_and_all_reduce( + [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue], + self.tp_worker.get_tp_group().cpu_group, + ) + + transferred_rids: List[str] = [] + + for req, poll in zip(self.disagg_prefill_inflight_queue, polls): + if poll == KVPoll.Success or poll == KVPoll.Failed: + transferred_rids.append(req.rid) + + return transferred_rids + + def process_prefill_chunk(self: Scheduler) -> None: + if self.last_batch and self.last_batch.forward_mode.is_extend(): + if self.chunked_req: + # Move the chunked request out of the batch so that we can merge + # only finished requests to running_batch. + self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req) + self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True) + if self.enable_overlap: + # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved + self.chunked_req.tmp_end_idx = min( + len(self.chunked_req.fill_ids), + len(self.chunked_req.origin_input_ids), + ) + else: + self.send_kv_chunk(self.chunked_req) + # chunked request keeps its rid but will get a new req_pool_idx + self.req_to_token_pool.free(self.chunked_req.req_pool_idx) + self.running_batch.batch_is_full = False + + def send_kv_chunk( + self: Scheduler, + req: Req, + last_chunk: bool = False, + end_idx: Optional[int] = None, + ) -> None: + """ + Send a prefilled chunk to the decode server + """ + page_size = self.token_to_kv_pool_allocator.page_size + start_idx = req.start_send_idx + end_idx = ( + end_idx + if end_idx is not None + else min(len(req.fill_ids), len(req.origin_input_ids)) + ) + + if not last_chunk: + # if not the last chunk and the last page is partial, delay the last partial page to the next send + end_idx = end_idx - end_idx % page_size + + kv_indices = ( + self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx] + .cpu() + .numpy() + ) + req.start_send_idx = end_idx + if last_chunk: + self.disagg_metadata_buffers.set_buf(req) + page_indices = kv_to_page_indices(kv_indices, page_size) + if len(page_indices) == 0: + logger.info( + f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty" + ) + return + req.disagg_kv_sender.send(page_indices) + + # PP + @DynamicGradMode() + def event_loop_pp_disagg_prefill(self: Scheduler): + """ + An event loop for the prefill server in pipeline parallelism. + + Rules: + 1. Each stage runs in the same order and is notified by the previous stage. + 2. Each send/recv operation is blocking and matched by the neighboring stage. + + Regular Schedule: + ==================================================================== + Stage i | Stage i+1 + send ith req | recv ith req + send ith proxy | recv ith proxy + send prev (i+1)th carry | recv prev (i+1)th carry + ==================================================================== + + Prefill Server Schedule: + ==================================================================== + Stage i | Stage i+1 + send ith req | recv ith req + send ith bootstrap req | recv ith bootstrap req + send ith transferred req | recv ith transferred req + send ith proxy | recv ith proxy + send prev (i+1)th carry | recv prev (i+1)th carry + send prev (i+1)th release req | recv prev (i+1)th release req + ==================================================================== + + There are two additional elements compared to the regular schedule: + + 1. Bootstrap Requests: + a. Instead of polling the status on the current workers, we should wait for the previous stage to notify to avoid desynchronization. + b. The first stage polls the status and propagates the bootstrapped requests down to all other stages. + c. If the first stage polls successfully, by nature, other ranks are also successful because they performed a handshake together. + + 2. Transferred Requests + Release Requests: + a. The first stage polls the transfer finished requests, performs an intersection with the next stage's finished requests, and propagates down to the last stage. + b. The last stage receives the requests that have finished transfer on all stages (consensus), then sends them to the first stage to release the memory. + c. The first stage receives the release requests, releases the memory, and then propagates the release requests down to the last stage. + """ + from sglang.srt.managers.scheduler import GenerationBatchResult + + mbs = [None] * self.pp_size + last_mbs = [None] * self.pp_size + self.running_mbs = [ + ScheduleBatch(reqs=[], batch_is_full=False) for _ in range(self.pp_size) + ] + bids = [None] * self.pp_size + pp_outputs: Optional[PPProxyTensors] = None + + # Either success or failed + bootstrapped_rids: List[str] = [] + transferred_rids: List[str] = [] + release_rids: Optional[List[str]] = None + + # transferred microbatch + tmbs = [None] * self.pp_size + + ENABLE_RELEASE = True # For debug + + while True: + server_is_idle = True + + for mb_id in range(self.pp_size): + self.running_batch = self.running_mbs[mb_id] + self.last_batch = last_mbs[mb_id] + + recv_reqs = self.recv_requests() + + self.process_input_requests(recv_reqs) + + if self.pp_group.is_first_rank: + # First rank, pop the bootstrap reqs from the bootstrap queue + bootstrapped_reqs, failed_reqs = ( + self.disagg_prefill_bootstrap_queue.pop_bootstrapped( + return_failed_reqs=True + ) + ) + bootstrapped_rids = [req.rid for req in bootstrapped_reqs] + [ + req.rid for req in failed_reqs + ] + self.waiting_queue.extend(bootstrapped_reqs) + else: + # Other ranks, receive the bootstrap reqs info from the previous rank and ensure the consensus + bootstrapped_rids = self.recv_pyobj_from_prev_stage() + bootstrapped_reqs = ( + self.disagg_prefill_bootstrap_queue.pop_bootstrapped( + rids_to_check=bootstrapped_rids + ) + ) + self.waiting_queue.extend(bootstrapped_reqs) + + if self.pp_group.is_first_rank: + transferred_rids = self.get_transferred_rids() + # if other ranks, + else: + # 1. recv previous stage's transferred reqs info + prev_transferred_rids = self.recv_pyobj_from_prev_stage() + # 2. get the current stage's transferred reqs info + curr_transferred_rids = self.get_transferred_rids() + # 3. new consensus rids = intersection(previous consensus rids, transfer finished rids) + transferred_rids = list( + set(prev_transferred_rids) & set(curr_transferred_rids) + ) + + tmbs[mb_id] = transferred_rids + + self.process_prefill_chunk() + mbs[mb_id] = self.get_new_batch_prefill() + self.running_mbs[mb_id] = self.running_batch + + self.cur_batch = mbs[mb_id] + if self.cur_batch: + server_is_idle = False + result = self.run_batch(self.cur_batch) + + # send the outputs to the next step + if self.pp_group.is_last_rank: + if self.cur_batch: + next_token_ids, bids[mb_id] = ( + result.next_token_ids, + result.bid, + ) + pp_outputs = PPProxyTensors( + { + "next_token_ids": next_token_ids, + } + ) + # send the output from the last round to let the next stage worker run post processing + self.pp_group.send_tensor_dict( + pp_outputs.tensors, + all_gather_group=self.attn_tp_group, + ) + + if ENABLE_RELEASE: + if self.pp_group.is_last_rank: + # At the last stage, all stages has reached the consensus to release memory for transferred_rids + release_rids = transferred_rids + # send to the first rank + self.send_pyobj_to_next_stage(release_rids) + + # receive outputs and post-process (filter finished reqs) the coming microbatch + next_mb_id = (mb_id + 1) % self.pp_size + next_pp_outputs = None + next_release_rids = None + + if mbs[next_mb_id] is not None: + next_pp_outputs: Optional[PPProxyTensors] = PPProxyTensors( + self.pp_group.recv_tensor_dict( + all_gather_group=self.attn_tp_group + ) + ) + mbs[next_mb_id].output_ids = next_pp_outputs["next_token_ids"] + output_result = GenerationBatchResult( + logits_output=None, + pp_hidden_states_proxy_tensors=None, + next_token_ids=next_pp_outputs["next_token_ids"], + extend_input_len_per_req=None, + extend_logprob_start_len_per_req=None, + bid=bids[next_mb_id], + can_run_cuda_graph=result.can_run_cuda_graph, + ) + self.process_batch_result_disagg_prefill( + mbs[next_mb_id], output_result + ) + + last_mbs[next_mb_id] = mbs[next_mb_id] + + if ENABLE_RELEASE: + if tmbs[next_mb_id] is not None: + # recv consensus rids from the previous rank + next_release_rids = self.recv_pyobj_from_prev_stage() + self.process_disagg_prefill_inflight_queue(next_release_rids) + + # carry the outputs to the next stage + if not self.pp_group.is_last_rank: + if self.cur_batch: + bids[mb_id] = result.bid + if pp_outputs: + # send the outputs from the last round to let the next stage worker run post processing + self.pp_group.send_tensor_dict( + pp_outputs.tensors, + all_gather_group=self.attn_tp_group, + ) + if ENABLE_RELEASE: + if release_rids is not None: + self.send_pyobj_to_next_stage(release_rids) + + if not self.pp_group.is_last_rank: + # send out reqs to the next stage + self.send_pyobj_to_next_stage(recv_reqs) + self.send_pyobj_to_next_stage(bootstrapped_rids) + self.send_pyobj_to_next_stage(transferred_rids) + + # send out proxy tensors to the next stage + if self.cur_batch: + self.pp_group.send_tensor_dict( + result.pp_hidden_states_proxy_tensors, + all_gather_group=self.attn_tp_group, + ) + + pp_outputs = next_pp_outputs + release_rids = next_release_rids + + self.running_batch.batch_is_full = False + + if not ENABLE_RELEASE: + if len(self.disagg_prefill_inflight_queue) > 0: + self.process_disagg_prefill_inflight_queue() + + # When the server is idle, self-check and re-init some states + if server_is_idle and len(self.disagg_prefill_inflight_queue) == 0: + self.check_memory() + self.check_tree_cache() + self.new_token_ratio = self.init_new_token_ratio + + def send_pyobj_to_next_stage(self, data): + if self.attn_tp_rank == 0: + dp_offset = self.attn_dp_rank * self.attn_tp_size + point_to_point_pyobj( + data, + self.pp_rank * self.tp_size + dp_offset, + self.world_group.device_group, + self.pp_rank * self.tp_size + dp_offset, + ((self.pp_rank + 1) % self.pp_size) * self.tp_size + dp_offset, + ) + + def recv_pyobj_from_prev_stage(self): + if self.attn_tp_rank == 0: + dp_offset = self.attn_dp_rank * self.attn_tp_size + data = point_to_point_pyobj( + [], + self.pp_rank * self.tp_size + dp_offset, + self.world_group.device_group, + ((self.pp_rank - 1) % self.pp_size) * self.tp_size + dp_offset, + self.pp_rank * self.tp_size + dp_offset, + ) + else: + data = None + + if self.tp_size != 1: + data = broadcast_pyobj( + data, self.tp_group.rank, self.tp_cpu_group, src=self.tp_group.ranks[0] + ) + return data diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py new file mode 100644 index 000000000..43770e3e2 --- /dev/null +++ b/python/sglang/srt/disaggregation/utils.py @@ -0,0 +1,329 @@ +from __future__ import annotations + +import os +import random +from collections import deque +from contextlib import nullcontext +from enum import Enum +from typing import TYPE_CHECKING, List, Optional, Type, Union + +import numpy as np +import torch +import torch.distributed as dist + +from sglang.srt.utils import is_npu + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + +######################### +# Constants & Enums +######################### +FAKE_BOOTSTRAP_HOST = "2.2.2.2" + + +class DisaggregationMode(Enum): + NULL = "null" + PREFILL = "prefill" + DECODE = "decode" + + +######################### +# Synchronization +######################### + +# env var for testing failure, convert to float explicitly +FAILURE_PROB = float(os.getenv("DISAGGREGATION_TEST_FAILURE_PROB", 0)) + + +def poll_and_all_reduce(pollers, gloo_group): + # at a certain prob, the poll is failed to simulate failure + if FAILURE_PROB > 0: + from sglang.srt.disaggregation.base import KVPoll + + polls = [ + int(KVPoll.Failed) if random.random() < FAILURE_PROB else int(poller.poll()) + for poller in pollers + ] + else: + polls = [int(poller.poll()) for poller in pollers] + tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu") + dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=gloo_group) + return tensor_to_reduce.tolist() + + +######################### +# Metadata Buffers +######################### + + +class ReqToMetadataIdxAllocator: + """A memory pool that maps a request to its first output token location.""" + + def __init__( + self, + size: int, + ): + self.size = size + self.free_slots = deque(list(range(size))) + + def available_size(self): + return len(self.free_slots) + + def alloc(self) -> Optional[int]: + if len(self.free_slots) == 0: + return None + + return self.free_slots.popleft() + + def free(self, free_index: int): + self.free_slots.append(free_index) + + +class MetadataBuffers: + def __init__( + self, + size: int, + hidden_size: int, + dtype: torch.dtype, + max_top_logprobs_num: int = 128, + custom_mem_pool: torch.cuda.MemPool = None, + ): + self.custom_mem_pool = custom_mem_pool + device = "cpu" + if is_npu(): + # For ascend backend, output tokens are placed in the NPU and will be transferred by D2D channel. + device = "npu" + elif self.custom_mem_pool: + # TODO(shangming): Fix me (use 'cuda') when nvlink_transport of Mooncake is bug-free + device = "cpu" + with ( + torch.cuda.use_mem_pool(self.custom_mem_pool) + if self.custom_mem_pool + else nullcontext() + ): + # TODO: abort top_logprobs_num > 128 in PD + + # We transfer the metadata of first output token to decode + # The minimal size for RDMA is 64Bytes, so we pad it to > 64Bytes + self.output_ids = torch.zeros((size, 16), dtype=torch.int32, device=device) + + self.output_token_logprobs_val = torch.zeros( + (size, 16), dtype=torch.float32, device=device + ) + self.output_token_logprobs_idx = torch.zeros( + (size, 16), dtype=torch.int32, device=device + ) + self.output_top_logprobs_val = torch.zeros( + (size, max_top_logprobs_num), dtype=torch.float32, device=device + ) + self.output_top_logprobs_idx = torch.zeros( + (size, max_top_logprobs_num), dtype=torch.int32, device=device + ) + self.output_hidden_states = torch.zeros( + (size, hidden_size), dtype=dtype, device=device + ) + + def get_buf_infos(self): + ptrs = [ + self.output_ids.data_ptr(), + self.output_token_logprobs_val.data_ptr(), + self.output_token_logprobs_idx.data_ptr(), + self.output_top_logprobs_val.data_ptr(), + self.output_top_logprobs_idx.data_ptr(), + self.output_hidden_states.data_ptr(), + ] + data_lens = [ + self.output_ids.nbytes, + self.output_token_logprobs_val.nbytes, + self.output_token_logprobs_idx.nbytes, + self.output_top_logprobs_val.nbytes, + self.output_top_logprobs_idx.nbytes, + self.output_hidden_states.nbytes, + ] + item_lens = [ + self.output_ids[0].nbytes, + self.output_token_logprobs_val[0].nbytes, + self.output_token_logprobs_idx[0].nbytes, + self.output_top_logprobs_val[0].nbytes, + self.output_top_logprobs_idx[0].nbytes, + self.output_hidden_states[0].nbytes, + ] + return ptrs, data_lens, item_lens + + def get_buf(self, idx: int): + return ( + self.output_ids[idx], + self.output_token_logprobs_val[idx], + self.output_token_logprobs_idx[idx], + self.output_top_logprobs_val[idx], + self.output_top_logprobs_idx[idx], + self.output_hidden_states[idx], + ) + + def set_buf(self, req: Req): + + self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0] + if req.return_logprob: + if req.output_token_logprobs_val: # not none or empty list + self.output_token_logprobs_val[req.metadata_buffer_index][0] = ( + req.output_token_logprobs_val[0] + ) + if req.output_token_logprobs_idx: # not none or empty list + self.output_token_logprobs_idx[req.metadata_buffer_index][0] = ( + req.output_token_logprobs_idx[0] + ) + + if req.output_top_logprobs_val: # not none or empty list + self.output_top_logprobs_val[req.metadata_buffer_index][ + : len(req.output_top_logprobs_val[0]) + ] = torch.tensor( + req.output_top_logprobs_val[0], dtype=torch.float32, device="cpu" + ) + if req.output_top_logprobs_idx: # not none or empty list + self.output_top_logprobs_idx[req.metadata_buffer_index][ + : len(req.output_top_logprobs_idx[0]) + ] = torch.tensor( + req.output_top_logprobs_idx[0], dtype=torch.int32, device="cpu" + ) + # for PD + spec decode + if req.hidden_states_tensor is not None: + self.output_hidden_states[req.metadata_buffer_index].copy_( + req.hidden_states_tensor + ) + + +######################### +# Transfer Backend +######################### + + +class TransferBackend(Enum): + MOONCAKE = "mooncake" + NIXL = "nixl" + ASCEND = "ascend" + FAKE = "fake" + + +class KVClassType(Enum): + KVARGS = "kvargs" + MANAGER = "manager" + SENDER = "sender" + RECEIVER = "receiver" + BOOTSTRAP_SERVER = "bootstrap_server" + + +def get_kv_class( + transfer_backend: TransferBackend, class_type: KVClassType +) -> Optional[Type]: + from sglang.srt.disaggregation.fake import FakeKVReceiver, FakeKVSender + + if transfer_backend == TransferBackend.MOONCAKE: + from sglang.srt.disaggregation.base import KVArgs + from sglang.srt.disaggregation.mooncake import ( + MooncakeKVBootstrapServer, + MooncakeKVManager, + MooncakeKVReceiver, + MooncakeKVSender, + ) + + class_mapping = { + KVClassType.KVARGS: KVArgs, + KVClassType.MANAGER: MooncakeKVManager, + KVClassType.SENDER: MooncakeKVSender, + KVClassType.RECEIVER: (MooncakeKVReceiver), + KVClassType.BOOTSTRAP_SERVER: MooncakeKVBootstrapServer, + } + return class_mapping.get(class_type) + elif transfer_backend == TransferBackend.ASCEND: + from sglang.srt.disaggregation.ascend import ( + AscendKVBootstrapServer, + AscendKVManager, + AscendKVReceiver, + AscendKVSender, + ) + from sglang.srt.disaggregation.base import KVArgs + + class_mapping = { + KVClassType.KVARGS: KVArgs, + KVClassType.MANAGER: AscendKVManager, + KVClassType.SENDER: AscendKVSender, + KVClassType.RECEIVER: (AscendKVReceiver), + KVClassType.BOOTSTRAP_SERVER: AscendKVBootstrapServer, + } + return class_mapping.get(class_type) + elif transfer_backend == TransferBackend.NIXL: + from sglang.srt.disaggregation.base import KVArgs + from sglang.srt.disaggregation.nixl import ( + NixlKVBootstrapServer, + NixlKVManager, + NixlKVReceiver, + NixlKVSender, + ) + + class_mapping = { + KVClassType.KVARGS: KVArgs, + KVClassType.MANAGER: NixlKVManager, + KVClassType.SENDER: NixlKVSender, + KVClassType.RECEIVER: (NixlKVReceiver), + KVClassType.BOOTSTRAP_SERVER: NixlKVBootstrapServer, + } + return class_mapping.get(class_type) + elif transfer_backend == TransferBackend.FAKE: + from sglang.srt.disaggregation.base import KVArgs + from sglang.srt.disaggregation.fake import FakeKVReceiver, FakeKVSender + + class_mapping = { + KVClassType.KVARGS: KVArgs, + KVClassType.SENDER: FakeKVSender, + KVClassType.RECEIVER: (FakeKVReceiver), + } + return class_mapping.get(class_type) + + raise ValueError(f"Unsupported transfer backend: {transfer_backend}") + + +######################### +# KV Pages +######################### + + +def kv_to_page_indices(kv_indices: np.ndarray, page_size: int): + # 1. The page is guaranteed to be full except the last page. + # 2. page index = kv_index // page_size + # The return vector is kv_indices[::page_size] // page_size + if page_size == 1: # shortcut + return kv_indices + + return kv_indices[::page_size] // page_size + + +def kv_to_page_num(num_kv_indices: int, page_size: int): + # ceil(num_kv_indices / page_size) + return (num_kv_indices + page_size - 1) // page_size + + +######################### +# Misc +######################### + + +def is_mla_backend(target_kv_pool) -> bool: + from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool + + return isinstance(target_kv_pool, MLATokenToKVPool) + + +def prepare_abort(req: Req, error_message: str, status_code=None): + from sglang.srt.managers.schedule_batch import FINISH_ABORT + + # populate finish metadata and stream output + req.finished_reason = FINISH_ABORT(error_message, status_code) + + if req.return_logprob: + req.input_token_logprobs_val = [] + req.input_token_logprobs_idx = [] + req.input_top_logprobs_val = [] + req.input_top_logprobs_idx = [] + req.input_token_ids_logprobs_val = [] + req.input_token_ids_logprobs_idx = [] diff --git a/python/sglang/srt/distributed/__init__.py b/python/sglang/srt/distributed/__init__.py new file mode 100644 index 000000000..12f802055 --- /dev/null +++ b/python/sglang/srt/distributed/__init__.py @@ -0,0 +1,3 @@ +from sglang.srt.distributed.communication_op import * +from sglang.srt.distributed.parallel_state import * +from sglang.srt.distributed.utils import * diff --git a/python/sglang/srt/distributed/communication_op.py b/python/sglang/srt/distributed/communication_op.py new file mode 100644 index 000000000..95600edfb --- /dev/null +++ b/python/sglang/srt/distributed/communication_op.py @@ -0,0 +1,35 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/communication_op.py + +from typing import Any, Dict, Optional, Union + +import torch +import torch.distributed + +from .parallel_state import get_tp_group + + +def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor: + """All-reduce the input tensor across model parallel group.""" + return get_tp_group().all_reduce(input_) + + +def tensor_model_parallel_all_gather( + input_: torch.Tensor, dim: int = -1 +) -> torch.Tensor: + """All-gather the input tensor across model parallel group.""" + return get_tp_group().all_gather(input_, dim) + + +def tensor_model_parallel_gather( + input_: torch.Tensor, dst: int = 0, dim: int = -1 +) -> Optional[torch.Tensor]: + """Gather the input tensor across model parallel group.""" + return get_tp_group().gather(input_, dst, dim) + + +def broadcast_tensor_dict( + tensor_dict: Optional[Dict[Any, Union[torch.Tensor, Any]]] = None, src: int = 0 +): + if not torch.distributed.is_initialized(): + return tensor_dict + return get_tp_group().broadcast_tensor_dict(tensor_dict, src) diff --git a/python/sglang/srt/distributed/device_communicators/cuda_wrapper.py b/python/sglang/srt/distributed/device_communicators/cuda_wrapper.py new file mode 100644 index 000000000..c902f3141 --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/cuda_wrapper.py @@ -0,0 +1,183 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/cuda_wrapper.py + +"""This file is a pure Python wrapper for the cudart library. +It avoids the need to compile a separate shared library, and is +convenient for use when we just need to call a few functions. +""" + +import ctypes +import logging +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +# this line makes it possible to directly load `libcudart.so` using `ctypes` +import torch # noqa + +logger = logging.getLogger(__name__) + +# === export types and functions from cudart to Python === +# for the original cudart definition, please check +# https://docs.nvidia.com/cuda/cuda-runtime-api/index.html + +cudaError_t = ctypes.c_int +cudaMemcpyKind = ctypes.c_int + + +class cudaIpcMemHandle_t(ctypes.Structure): + _fields_ = [("internal", ctypes.c_byte * 128)] + + +@dataclass +class Function: + name: str + restype: Any + argtypes: List[Any] + + +def find_loaded_library(lib_name) -> Optional[str]: + """ + According to according to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html, + the file `/proc/self/maps` contains the memory maps of the process, which includes the + shared libraries loaded by the process. We can use this file to find the path of the + a loaded library. + """ # noqa + found = False + with open("/proc/self/maps") as f: + for line in f: + if lib_name in line: + found = True + break + if not found: + # the library is not loaded in the current process + return None + # if lib_name is libcudart, we need to match a line with: + # address /path/to/libcudart-hash.so.11.0 + start = line.index("/") + path = line[start:].strip() + filename = path.split("/")[-1] + assert filename.rpartition(".so")[0].startswith( + lib_name + ), f"Unexpected filename: {filename} for library {lib_name}" + return path + + +class CudaRTLibrary: + exported_functions = [ + # ​cudaError_t cudaSetDevice ( int device ) + Function("cudaSetDevice", cudaError_t, [ctypes.c_int]), + # cudaError_t cudaDeviceSynchronize ( void ) + Function("cudaDeviceSynchronize", cudaError_t, []), + # ​cudaError_t cudaDeviceReset ( void ) + Function("cudaDeviceReset", cudaError_t, []), + # const char* cudaGetErrorString ( cudaError_t error ) + Function("cudaGetErrorString", ctypes.c_char_p, [cudaError_t]), + # ​cudaError_t cudaMalloc ( void** devPtr, size_t size ) + Function( + "cudaMalloc", + cudaError_t, + [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t], + ), + # ​cudaError_t cudaFree ( void* devPtr ) + Function("cudaFree", cudaError_t, [ctypes.c_void_p]), + # ​cudaError_t cudaMemset ( void* devPtr, int value, size_t count ) + Function( + "cudaMemset", cudaError_t, [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t] + ), + # ​cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind ) # noqa + Function( + "cudaMemcpy", + cudaError_t, + [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, cudaMemcpyKind], + ), + # cudaError_t cudaIpcGetMemHandle ( cudaIpcMemHandle_t* handle, void* devPtr ) # noqa + Function( + "cudaIpcGetMemHandle", + cudaError_t, + [ctypes.POINTER(cudaIpcMemHandle_t), ctypes.c_void_p], + ), + # ​cudaError_t cudaIpcOpenMemHandle ( void** devPtr, cudaIpcMemHandle_t handle, unsigned int flags ) # noqa + Function( + "cudaIpcOpenMemHandle", + cudaError_t, + [ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint], + ), + ] + + # class attribute to store the mapping from the path to the library + # to avoid loading the same library multiple times + path_to_library_cache: Dict[str, Any] = {} + + # class attribute to store the mapping from library path + # to the corresponding dictionary + path_to_dict_mapping: Dict[str, Dict[str, Any]] = {} + + def __init__(self, so_file: Optional[str] = None): + if so_file is None: + so_file = find_loaded_library("libcudart") + assert so_file is not None, "libcudart is not loaded in the current process" + if so_file not in CudaRTLibrary.path_to_library_cache: + lib = ctypes.CDLL(so_file) + CudaRTLibrary.path_to_library_cache[so_file] = lib + self.lib = CudaRTLibrary.path_to_library_cache[so_file] + + if so_file not in CudaRTLibrary.path_to_dict_mapping: + _funcs = {} + for func in CudaRTLibrary.exported_functions: + f = getattr(self.lib, func.name) + f.restype = func.restype + f.argtypes = func.argtypes + _funcs[func.name] = f + CudaRTLibrary.path_to_dict_mapping[so_file] = _funcs + self.funcs = CudaRTLibrary.path_to_dict_mapping[so_file] + + def CUDART_CHECK(self, result: cudaError_t) -> None: + if result != 0: + error_str = self.cudaGetErrorString(result) + raise RuntimeError(f"CUDART error: {error_str}") + + def cudaGetErrorString(self, error: cudaError_t) -> str: + return self.funcs["cudaGetErrorString"](error).decode("utf-8") + + def cudaSetDevice(self, device: int) -> None: + self.CUDART_CHECK(self.funcs["cudaSetDevice"](device)) + + def cudaDeviceSynchronize(self) -> None: + self.CUDART_CHECK(self.funcs["cudaDeviceSynchronize"]()) + + def cudaDeviceReset(self) -> None: + self.CUDART_CHECK(self.funcs["cudaDeviceReset"]()) + + def cudaMalloc(self, size: int) -> ctypes.c_void_p: + devPtr = ctypes.c_void_p() + self.CUDART_CHECK(self.funcs["cudaMalloc"](ctypes.byref(devPtr), size)) + return devPtr + + def cudaFree(self, devPtr: ctypes.c_void_p) -> None: + self.CUDART_CHECK(self.funcs["cudaFree"](devPtr)) + + def cudaMemset(self, devPtr: ctypes.c_void_p, value: int, count: int) -> None: + self.CUDART_CHECK(self.funcs["cudaMemset"](devPtr, value, count)) + + def cudaMemcpy( + self, dst: ctypes.c_void_p, src: ctypes.c_void_p, count: int + ) -> None: + cudaMemcpyDefault = 4 + kind = cudaMemcpyDefault + self.CUDART_CHECK(self.funcs["cudaMemcpy"](dst, src, count, kind)) + + def cudaIpcGetMemHandle(self, devPtr: ctypes.c_void_p) -> cudaIpcMemHandle_t: + handle = cudaIpcMemHandle_t() + self.CUDART_CHECK( + self.funcs["cudaIpcGetMemHandle"](ctypes.byref(handle), devPtr) + ) + return handle + + def cudaIpcOpenMemHandle(self, handle: cudaIpcMemHandle_t) -> ctypes.c_void_p: + cudaIpcMemLazyEnablePeerAccess = 1 + devPtr = ctypes.c_void_p() + self.CUDART_CHECK( + self.funcs["cudaIpcOpenMemHandle"]( + ctypes.byref(devPtr), handle, cudaIpcMemLazyEnablePeerAccess + ) + ) + return devPtr diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py new file mode 100644 index 000000000..0a506d35f --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -0,0 +1,421 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/custom_all_reduce.py + +import ctypes +import logging +import os +from contextlib import contextmanager +from typing import Any, List, Optional, Union + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from sglang.srt import _custom_ops as ops +from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary +from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import ( + gpu_p2p_access_check, + is_full_nvlink, + is_weak_contiguous, +) +from sglang.srt.distributed.parallel_state import in_the_same_node_as +from sglang.srt.utils import is_cuda, is_hip + +logger = logging.getLogger(__name__) + +_is_cuda = is_cuda() +_is_hip = is_hip() + + +try: + if ops.use_vllm_custom_allreduce and not _is_hip: + # Use vLLM custom allreduce + ops.meta_size() + else: + # Use custom allreduce from sgl kernel (ROCM and TRT-LLM) + import sgl_kernel + custom_ar = True +except Exception: + # For CPUs + custom_ar = False + +logger = logging.getLogger(__name__) + + +def _can_p2p(rank: int, world_size: int) -> bool: + # SGLANG_SKIP_P2P_CHECK can be set to False in sglang + SGLANG_SKIP_P2P_CHECK = os.getenv("SGLANG_SKIP_P2P_CHECK", "0") == "1" + for i in range(world_size): + if i == rank: + continue + if SGLANG_SKIP_P2P_CHECK: + logger.info("Skipping P2P check and trusting the driver's P2P report.") + return torch.cuda.can_device_access_peer(rank, i) + if not gpu_p2p_access_check(rank, i): + return False + return True + + +class CustomAllreduce: + _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8] + _MAX_CAR_SIZE = 8192 * 1024 + if _is_hip: + # crossover is at 16MB buffer size for ROCm + _MAX_CAR_SIZE = 2 * 8192 * 1024 + + # max_size: max supported allreduce size + def __init__( + self, + group: ProcessGroup, + device: Union[int, str, torch.device], + max_size=_MAX_CAR_SIZE, + ) -> None: + """ + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the CustomAllreduce to. If None, + it will be bind to f"cuda:{local_rank}". + It is the caller's responsibility to make sure each communicator + is bind to a unique device, and all communicators in this group + are in the same node. + """ + self._IS_CAPTURING = False + self.disabled = True + + if not custom_ar: + # disable because of missing custom allreduce library + # e.g. in a non-cuda environment + return + + self.group = group + + assert ( + dist.get_backend(group) != dist.Backend.NCCL + ), "CustomAllreduce should be attached to a non-NCCL group." + + if not all(in_the_same_node_as(group, source_rank=0)): + # No need to initialize custom allreduce for multi-node case. + logger.warning( + "Custom allreduce is disabled because this process group" + " spans across nodes." + ) + return + + rank = dist.get_rank(group=self.group) + world_size = dist.get_world_size(group=self.group) + if world_size == 1: + # No need to initialize custom allreduce for single GPU case. + return + + if world_size not in CustomAllreduce._SUPPORTED_WORLD_SIZES: + logger.warning( + "Custom allreduce is disabled due to an unsupported world" + " size: %d. Supported world sizes: %s. To silence this " + "warning, specify disable_custom_all_reduce=True explicitly.", + world_size, + str(CustomAllreduce._SUPPORTED_WORLD_SIZES), + ) + return + + if isinstance(device, int): + device = torch.device(f"cuda:{device}") + elif isinstance(device, str): + device = torch.device(device) + # now `device` is a `torch.device` object + assert isinstance(device, torch.device) + self.device = device + + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None) + if cuda_visible_devices: + device_ids = list(map(int, cuda_visible_devices.split(","))) + else: + device_ids = list(range(torch.cuda.device_count())) + + physical_device_id = device_ids[device.index] + tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") + gather_list = [ + torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size) + ] + dist.all_gather(gather_list, tensor, group=self.group) + physical_device_ids = [t.item() for t in gather_list] + + # test nvlink first, this will filter out most of the cases + # where custom allreduce is not supported + # this checks hardware and driver support for NVLink + if _is_cuda or _is_hip: + full_nvlink = is_full_nvlink(physical_device_ids, world_size) + + if world_size > 2 and not full_nvlink: + logger.warning( + "Custom allreduce is disabled because it's not supported on" + " more than two PCIe-only GPUs. To silence this warning, " + "specify disable_custom_all_reduce=True explicitly." + ) + return + # test P2P capability, this checks software/cudaruntime support + # this is expensive to compute at the first time + # then we cache the result + # On AMD GPU, p2p is always enabled between XGMI connected GPUs + if not _is_hip and not _can_p2p(rank, world_size): + logger.warning( + "Custom allreduce is disabled because your platform lacks " + "GPU P2P capability or P2P test failed. To silence this " + "warning, specify disable_custom_all_reduce=True explicitly." + ) + return + + self.max_size = max_size + self.rank = rank + self.world_size = world_size + self.full_nvlink = full_nvlink + + if not _is_hip: + # Buffers memory are owned by this Python class and passed to C++. + # Meta data composes of two parts: meta data for synchronization and a + # temporary buffer for storing intermediate allreduce results. + self.meta_ptrs = self.create_shared_buffer( + ops.meta_size() + max_size, group=group + ) + # This is a pre-registered IPC buffer. In eager mode, input tensors + # are first copied into this buffer before allreduce is performed + self.buffer_ptrs = self.create_shared_buffer(max_size, group=group) + # This is a buffer for storing the tuples of pointers pointing to + # IPC buffers from all ranks. Each registered tuple has size of + # 8*world_size bytes where world_size is at most 8. Allocating 8MB + # is enough for 131072 such tuples. The largest model I've seen only + # needs less than 10000 of registered tuples. + self.rank_data = torch.empty( + 8 * 1024 * 1024, dtype=torch.uint8, device=self.device + ) + self._ptr = ops.init_custom_ar( + self.meta_ptrs, self.rank_data, rank, self.full_nvlink + ) + ops.register_buffer(self._ptr, self.buffer_ptrs) + else: + # meta data buffers need to be "uncached" for signal on MI200 + self.meta = ops.allocate_meta_buffer(ops.meta_size() + max_size) + self.buffer = torch.empty(max_size, dtype=torch.uint8, device=self.device) + handle = ops.get_meta_buffer_ipc_handle(self.meta) + shard_data = ( + bytes(handle), # ipc handle to base ptr + 0, # offset of base ptr + ) + handles, offsets = self._gather_ipc_meta(shard_data) + self.rank_data = torch.empty( + 8 * 1024 * 1024, dtype=torch.uint8, device=self.device + ) + self._ptr = ops.init_custom_ar( + self.meta, self.rank_data, handles, offsets, rank, self.full_nvlink + ) + self.register_buffer(self.buffer) + + self.disabled = False + + @staticmethod + def create_shared_buffer( + size_in_bytes: int, group: Optional[ProcessGroup] = None + ) -> List[int]: + """ + Creates a shared buffer and returns a list of pointers + representing the buffer on all processes in the group. + """ + lib = CudaRTLibrary() + pointer = lib.cudaMalloc(size_in_bytes) + handle = lib.cudaIpcGetMemHandle(pointer) + world_size = dist.get_world_size(group=group) + rank = dist.get_rank(group=group) + handles = [None] * world_size + dist.all_gather_object(handles, handle, group=group) + + pointers: List[int] = [] + for i, h in enumerate(handles): + if i == rank: + pointers.append(pointer.value) # type: ignore + else: + pointers.append(lib.cudaIpcOpenMemHandle(h).value) # type: ignore + + return pointers + + @staticmethod + def free_shared_buffer( + pointers: List[int], group: Optional[ProcessGroup] = None + ) -> None: + rank = dist.get_rank(group=group) + lib = CudaRTLibrary() + lib.cudaFree(ctypes.c_void_p(pointers[rank])) + + @contextmanager + def capture(self): + """ + The main responsibility of this context manager is the + `register_graph_buffers` call at the end of the context. + It records all the buffer addresses used in the CUDA graph. + """ + try: + self._IS_CAPTURING = True + yield + finally: + self._IS_CAPTURING = False + if not self.disabled: + self.register_graph_buffers() + + def _get_ipc_meta(self, inp: torch.Tensor): + # _share_cuda_() doesn't accept meta buffer not allocated from + # PyTorch cache allocator, use direct HIP call to get IPC handle + handle = ops.get_meta_buffer_ipc_handle(inp) + shard_data = ( + bytes(handle), # ipc handle to base ptr + 0, # offset of base ptr + ) + return self._gather_ipc_meta(shard_data) + + def _gather_ipc_meta(self, shard_data): + # Note: don't use `[[None]] * self.world_size` here + # because it will create a list of the same reference + all_data: List[Optional[Any]] = [[None] for i in range(self.world_size)] + all_data[self.rank][0] = shard_data + + ranks = dist.get_process_group_ranks(group=self.group) + ranks.sort() + for i, rank in enumerate(ranks): + dist.broadcast_object_list( + all_data[i], src=rank, group=self.group, device="cpu" + ) + + # we cannot directly use `dist.all_gather_object` here + # because it is incompatible with `gloo` backend under inference mode. + # see https://github.com/pytorch/pytorch/issues/126032 for details. + + handles = [] + offsets = [] + for i in range(len(all_data)): + handles.append(all_data[i][0][0]) # type: ignore + offsets.append(all_data[i][0][1]) # type: ignore + return handles, offsets + + def register_buffer(self, inp: torch.Tensor): + handles, offsets = self._get_ipc_meta(inp) + ops.register_buffer(self._ptr, inp, handles, offsets) + + def register_graph_buffers(self): + if _is_hip: + handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr) + handles, offsets = self._gather_ipc_meta((bytes(handle), offset)) + logger.info("Registering %d cuda graph addresses", len(offset)) + ops.register_graph_buffers(self._ptr, handles, offsets) + else: + handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr) + logger.info("Registering %d cuda graph addresses", len(offset)) + # We cannot directly use `dist.all_gather_object` here + # because it is incompatible with `gloo` backend under inference mode. + # see https://github.com/pytorch/pytorch/issues/126032 for details. + all_data = [ + [None, None] for _ in range(dist.get_world_size(group=self.group)) + ] + all_data[self.rank] = [handle, offset] + ranks = sorted(dist.get_process_group_ranks(group=self.group)) + for i, rank in enumerate(ranks): + dist.broadcast_object_list( + all_data[i], src=rank, group=self.group, device="cpu" + ) + # Unpack list of tuples to tuple of lists. + handles = [d[0] for d in all_data] # type: ignore + offsets = [d[1] for d in all_data] # type: ignore + ops.register_graph_buffers(self._ptr, handles, offsets) + + def should_custom_ar(self, inp: torch.Tensor): + if self.disabled: + return False + inp_size = inp.numel() * inp.element_size() + # custom allreduce requires input byte size to be multiples of 16 + if inp_size % 16 != 0: + return False + if not is_weak_contiguous(inp): + return False + # for 4 or more non NVLink-capable GPUs, custom allreduce provides + # little performance improvement over NCCL. + if not _is_hip: + if self.world_size == 2 or self.full_nvlink: + return inp_size < self.max_size + return False + + if _is_hip: + if self.full_nvlink: + return inp_size < self.max_size + return False + + return False + + # all reduce, assuming inp tensor is IPC registered with register_buffer, + # or, in the context of cuda graphs, register_graph_buffers + def all_reduce_reg(self, inp: torch.Tensor, out: torch.Tensor = None): + if out is None: + out = torch.empty_like(inp) + ops.all_reduce_reg(self._ptr, inp, out) + return out + + # all reduce, assuming inp tensor is NOT IPC registered + def all_reduce_unreg(self, inp: torch.Tensor, out: torch.Tensor = None): + if out is None: + out = torch.empty_like(inp) + ops.all_reduce_unreg(self._ptr, inp, self.buffer, out) + return out + + def all_reduce( + self, + inp: torch.Tensor, + *, + out: torch.Tensor = None, + registered: bool = False, + ): + """Performs an out-of-place all reduce. + + If registered is True, this assumes inp's pointer is already + IPC-registered. Otherwise, inp is first copied into a pre-registered + buffer. + """ + if out is None: + out = torch.empty_like(inp) + if registered: + ops.all_reduce(self._ptr, inp, out, 0, 0) + else: + ops.all_reduce( + self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size + ) + return out + + def custom_all_reduce(self, input: torch.Tensor) -> Optional[torch.Tensor]: + """The main allreduce API that provides support for cuda graph.""" + # When custom allreduce is disabled, this will be None. + if self.disabled or not self.should_custom_ar(input): + return None + if self._IS_CAPTURING: + if torch.cuda.is_current_stream_capturing(): + if _is_hip: + return self.all_reduce_reg(input) + else: + return self.all_reduce(input, registered=True) + else: + # If warm up, mimic the allocation pattern since custom + # allreduce is out-of-place. + return torch.zeros_like(input) + else: + if _is_hip: + # note: outside of cuda graph context, + # custom allreduce incurs a cost of cudaMemcpy, which should + # be small(<=1% of overall latency) compared to the performance + # gains of using custom kernels + return self.all_reduce_unreg(input) + else: + return self.all_reduce(input, registered=False) + + def close(self): + if not self.disabled and self._ptr: + ops.dispose(self._ptr) + if _is_cuda: + self.free_shared_buffer(self.meta_ptrs) + self.free_shared_buffer(self.buffer_ptrs) + self._ptr = 0 + + def __del__(self): + self.close() diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py new file mode 100644 index 000000000..c7baac845 --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py @@ -0,0 +1,386 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/custom_all_reduce_utils.py + +import ctypes +import json +import logging +import os +import pickle +import subprocess +import sys +import tempfile +from functools import wraps +from itertools import product +from typing import Callable, Dict, List, Optional, Sequence, TypeVar + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from typing_extensions import ParamSpec + +from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary +from sglang.srt.utils import is_cuda, is_hip + +logger = logging.getLogger(__name__) + +_is_cuda = is_cuda() +_is_hip = is_hip() + +if _is_cuda: + try: + import pynvml + except ImportError as e: + logger.warning("Failed to import pynvml with %r", e) + +if _is_hip: + try: + from amdsmi import ( + AmdSmiException, + amdsmi_get_processor_handles, + amdsmi_init, + amdsmi_shut_down, + amdsmi_topo_get_link_type, + ) + except ImportError as e: + logger.warning("Failed to import amdsmi with %r", e) + +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +def update_environment_variables(envs: Dict[str, str]): + for k, v in envs.items(): + if k in os.environ and os.environ[k] != v: + logger.warning( + "Overwriting environment variable %s " "from '%s' to '%s'", + k, + os.environ[k], + v, + ) + os.environ[k] = v + + +def producer( + batch_src: Sequence[int], + producer_queue, + consumer_queue, + result_queue, + cuda_visible_devices: Optional[str] = None, +): + if cuda_visible_devices is not None: + update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices}) + + lib = CudaRTLibrary() + for i in batch_src: + lib.cudaSetDevice(i) + pointer = lib.cudaMalloc(1024) + lib.cudaMemset(pointer, 1, 1024) + lib.cudaDeviceSynchronize() + handle = lib.cudaIpcGetMemHandle(pointer) + producer_queue.put(handle) + open_success = consumer_queue.get() + if open_success: + # use two queues to simulate barrier + producer_queue.put(0) + consumer_queue.get() + # check if the memory is modified + host_data = (ctypes.c_char * 1024)() + lib.cudaMemcpy(host_data, pointer, 1024) # type: ignore + for i in range(1024): + if ord(host_data[i]) != 2: + open_success = False + break + result_queue.put(open_success) + lib.cudaDeviceReset() + + +def consumer( + batch_tgt: Sequence[int], + producer_queue, + consumer_queue, + result_queue, + cuda_visible_devices: Optional[str] = None, +): + if cuda_visible_devices is not None: + update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices}) + + lib = CudaRTLibrary() + for j in batch_tgt: + lib.cudaSetDevice(j) + handle = producer_queue.get() + open_success = False + try: + pointer = lib.cudaIpcOpenMemHandle(handle) # type: ignore + open_success = True + except RuntimeError: + # cannot error out here, because the producer process + # is still waiting for the response. + pass + consumer_queue.put(open_success) + if open_success: + # modify the memory + lib.cudaMemset(pointer, 2, 1024) + lib.cudaDeviceSynchronize() + # use two queues to simulate barrier + producer_queue.get() + consumer_queue.put(0) + # check if the memory is modified + host_data = (ctypes.c_char * 1024)() + lib.cudaMemcpy(host_data, pointer, 1024) # type: ignore + for i in range(1024): + if ord(host_data[i]) != 2: + open_success = False + break + result_queue.put(open_success) + lib.cudaDeviceReset() + + +def can_actually_p2p( + batch_src: Sequence[int], + batch_tgt: Sequence[int], +) -> Sequence[bool]: + """ + Usually, checking if P2P access is enabled can be done by + `torch.cuda.can_device_access_peer(src, tgt)`. However, sometimes + the driver might be broken, and `torch.cuda.can_device_access_peer(src, tgt)` + returns `True` even if P2P access is not actually possible. + See https://github.com/vllm-project/vllm/issues/2728 and + https://forums.developer.nvidia.com/t/direct-gpu-gpu-communication-does-not-seem-to-work-properly/283264/10 + Therefore, we have to perform a real P2P access to check if it is actually + possible. + + Note on p2p and cuda IPC: + Usually, one process uses one GPU: + GPU src --> cuda context src --> tensor src --> process src + + We need to combine p2p and cuda IPC, so that: + GPU src --> cuda context src --> tensor src --> process src + |shared| + GPU tgt --> cuda context tgt --> tensor tgt --> process tgt + That is to say, process src creates a tensor in GPU src, passes IPC handle to + process tgt, and process tgt accesses the tensor in GPU tgt. Any operation on the + tensor in process tgt will be reflected in the tensor in process src, because + they are the same memory segment. + It is important to note that process tgt accesses the tensor in GPU tgt, not + GPU src. That's why we need p2p access. + + The most time-consuming part is the process creation. To avoid creating + processes for every pair of GPUs, we use batched testing. We create two + processes for testing all pairs of GPUs in batch. The trick is to reset + the device after each test (which is not available in PyTorch). + """ # noqa + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None) + # pass the CUDA_VISIBLE_DEVICES to the child process + # to make sure they see the same set of GPUs + + # make sure the processes are spawned + smp = mp.get_context("spawn") + producer_queue = smp.Queue() + consumer_queue = smp.Queue() + result_queue = smp.Queue() + p_src = smp.Process( + target=producer, + args=( + batch_src, + producer_queue, + consumer_queue, + result_queue, + cuda_visible_devices, + ), + ) + p_tgt = smp.Process( + target=consumer, + args=( + batch_tgt, + producer_queue, + consumer_queue, + result_queue, + cuda_visible_devices, + ), + ) + p_src.start() + p_tgt.start() + p_src.join() + p_tgt.join() + assert p_src.exitcode == 0 and p_tgt.exitcode == 0 + result: List[bool] = [] + for src, tgt in zip(batch_src, batch_tgt): + a = result_queue.get() + b = result_queue.get() + if a != b: + logger.warning( + "Two processes do not agree on the P2P access" + " status on %d -> %d, treat as disabled.", + src, + tgt, + ) + result.append(False) + else: + result.append(a) + return result + + +# why do we need this cache? +# we are testing peer-to-peer (p2p) access between GPUs,across processes. +# if we test it every time, it will be very slow, because we need to create +# N * N * 2 processes, where N is the world size. This is very slow. +# to reduce the time, we use a cache file to store the p2p access status. +# the cache file is generated by the master process if it does not exist. +# then all the processes can read the cache file to check the p2p access status. +# Note that the cache file is suffixed by the CUDA_VISIBLE_DEVICES, so that we +# can have different cache files for different CUDA_VISIBLE_DEVICES settings, +# e.g. used by different vllm engines. The device id in the cache file is a +# **local** device id, i.e. from 0 to num_dev-1, where num_dev is the number +# of visible devices in the vllm engine. +_gpu_p2p_access_cache: Optional[Dict[str, bool]] = None + + +def gpu_p2p_access_check(src: int, tgt: int) -> bool: + """Check if GPU src can access GPU tgt.""" + + # if the cache variable is already calculated, + # read from the cache instead of checking it again + global _gpu_p2p_access_cache + if _gpu_p2p_access_cache is not None: + return _gpu_p2p_access_cache[f"{src}->{tgt}"] + + is_distributed = dist.is_initialized() + + num_dev = torch.cuda.device_count() + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None) + if cuda_visible_devices is None: + cuda_visible_devices = ",".join(str(i) for i in range(num_dev)) + + # VLLM_CACHE_ROOT -> SGLANG_CACHE_ROOT + # "~/.cache/vllm" -> "~/.cache/sglang" + SGLANG_CACHE_ROOT = os.path.expanduser("~/.cache/sglang") + path = os.path.join( + SGLANG_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json" + ) + os.makedirs(os.path.dirname(path), exist_ok=True) + from sglang.srt.distributed.parallel_state import get_world_group + + if (not is_distributed or get_world_group().local_rank == 0) and ( + not os.path.exists(path) + ): + # only the local master process (with local_rank == 0) can + # enter this block to calculate the cache + logger.info("generating GPU P2P access cache in %s", path) + cache: Dict[str, bool] = {} + ids = list(range(num_dev)) + # batch of all pairs of GPUs + batch_src, batch_tgt = zip(*list(product(ids, ids))) + # NOTE: we use `subprocess` rather than `multiprocessing` here + # because the caller might not have `if __name__ == "__main__":`, + # in that case we cannot use spawn method in multiprocessing. + # However, `can_actually_p2p` requires spawn method. + # The fix is, we use `subprocess` to call the function, + # where we have `if __name__ == "__main__":` in this file. + + # use a temporary file to store the result + # we don't use the output of the subprocess directly, + # because the subprocess might produce logging output + with tempfile.NamedTemporaryFile() as output_file: + input_bytes = pickle.dumps((batch_src, batch_tgt, output_file.name)) + returned = subprocess.run( + [sys.executable, __file__], input=input_bytes, capture_output=True + ) + # check if the subprocess is successful + try: + returned.check_returncode() + except Exception as e: + # wrap raised exception to provide more information + raise RuntimeError( + f"Error happened when batch testing " + f"peer-to-peer access from {batch_src} to {batch_tgt}:\n" + f"{returned.stderr.decode()}" + ) from e + with open(output_file.name, "rb") as f: + result = pickle.load(f) + for _i, _j, r in zip(batch_src, batch_tgt, result): + cache[f"{_i}->{_j}"] = r + with open(path, "w") as f: + json.dump(cache, f, indent=4) + if is_distributed: + get_world_group().barrier() + logger.info("reading GPU P2P access cache from %s", path) + with open(path) as f: + cache = json.load(f) + _gpu_p2p_access_cache = cache + return _gpu_p2p_access_cache[f"{src}->{tgt}"] + + +def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]: + @wraps(fn) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + if _is_hip: + try: + amdsmi_init() + return fn(*args, **kwargs) + finally: + amdsmi_shut_down() + else: + pynvml.nvmlInit() + try: + return fn(*args, **kwargs) + finally: + pynvml.nvmlShutdown() + + return wrapper + + +@with_nvml_context +def is_full_nvlink(physical_device_ids: List[int], world_size: int) -> bool: + if _is_hip: + """ + query if the set of gpus are fully connected by xgmi (1 hop) + """ + handles = [amdsmi_get_processor_handles()[i] for i in physical_device_ids] + for i, handle in enumerate(handles): + for j, peer_handle in enumerate(handles): + if i < j: + try: + link_type = amdsmi_topo_get_link_type(handle, peer_handle) + # type is 2 for XGMI + if link_type["hops"] != 1 or link_type["type"] != 2: + return False + except AmdSmiException as error: + logger.error("AMD 1 hop XGMI detection failed.", exc_info=error) + return False + return True + else: + """ + query if the set of gpus are fully connected by nvlink (1 hop) + """ + handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids] + for i, handle in enumerate(handles): + for j, peer_handle in enumerate(handles): + if i < j: + try: + p2p_status = pynvml.nvmlDeviceGetP2PStatus( + handle, peer_handle, pynvml.NVML_P2P_CAPS_INDEX_NVLINK + ) + if p2p_status != pynvml.NVML_P2P_STATUS_OK: + return False + except pynvml.NVMLError: + logger.exception( + "NVLink detection failed. This is normal if your" + " machine has no NVLink equipped." + ) + return False + return True + + +def is_weak_contiguous(inp: torch.Tensor): + return inp.is_contiguous() or ( + inp.storage().nbytes() - inp.storage_offset() * inp.element_size() + == inp.numel() * inp.element_size() + ) + + +__all__ = ["gpu_p2p_access_check"] + +if __name__ == "__main__": + batch_src, batch_tgt, output_file = pickle.loads(sys.stdin.buffer.read()) + result = can_actually_p2p(batch_src, batch_tgt) + with open(output_file, "wb") as f: + f.write(pickle.dumps(result)) diff --git a/python/sglang/srt/distributed/device_communicators/hpu_communicator.py b/python/sglang/srt/distributed/device_communicators/hpu_communicator.py new file mode 100644 index 000000000..722e494cf --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/hpu_communicator.py @@ -0,0 +1,49 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/hpu_communicator.py + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from sglang.srt.utils import is_hpu + +if is_hpu(): + import habana_frameworks.torch as htorch # noqa: F401 + + +class HpuCommunicator: + + def __init__(self, group: ProcessGroup): + if not is_hpu(): + self.disabled = True + return + self.disabled = False + self.group = group + self.world_size = dist.get_world_size(self.group) + + def all_reduce(self, x: torch.Tensor) -> torch.Tensor: + # FIXME(kzawora): this is a workaround for a bug in Habana PT bridge + # occurring when PT_HPU_ENABLE_LAZY_COLLECTIVES=true env var is used + # (which is required for tensor parallel HPUGraph inference) + htorch.core.mark_step() + dist.all_reduce(x, group=self.group) + return x + + def all_gather(self, x: torch.Tensor, dim: int = -1) -> torch.Tensor: + world_size = self.world_size + if dim < 0: + # Convert negative dim to positive. + dim += x.dim() + input_size = x.size() + # Allocate output tensor. + output_tensor = torch.empty( + (world_size,) + input_size, dtype=x.dtype, device=x.device + ) + # All-gather. + htorch.core.mark_step() + dist.all_gather_into_tensor(output_tensor, x, group=self.group) + # Reshape + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :] + ) + return output_tensor diff --git a/python/sglang/srt/distributed/device_communicators/npu_communicator.py b/python/sglang/srt/distributed/device_communicators/npu_communicator.py new file mode 100644 index 000000000..cb6eb88e3 --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/npu_communicator.py @@ -0,0 +1,39 @@ +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from sglang.srt.utils import is_npu + + +class NpuCommunicator: + + def __init__(self, group: ProcessGroup): + if not is_npu(): + self.disabled = True + return + self.disabled = False + self.group = group + self.world_size = dist.get_world_size(self.group) + + def all_reduce(self, x: torch.Tensor) -> torch.Tensor: + dist.all_reduce(x, group=self.group) + return x + + def all_gather(self, x: torch.Tensor, dim: int = -1) -> torch.Tensor: + world_size = self.world_size + if dim < 0: + # Convert negative dim to positive. + dim += x.dim() + input_size = x.size() + output_size = (input_size[0] * world_size,) + input_size[1:] + # Allocate output tensor. + output_tensor = torch.empty(output_size, dtype=x.dtype, device=x.device) + # All-gather. + dist.all_gather_into_tensor(output_tensor, x, group=self.group) + # Reshape + output_tensor = output_tensor.reshape((world_size,) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :] + ) + return output_tensor diff --git a/python/sglang/srt/distributed/device_communicators/pymscclpp.py b/python/sglang/srt/distributed/device_communicators/pymscclpp.py new file mode 100644 index 000000000..78269ed05 --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/pymscclpp.py @@ -0,0 +1,315 @@ +import bisect +import logging +import math +import os +from contextlib import contextmanager +from enum import IntEnum +from typing import Any, Callable, List, Optional, TypeVar, Union + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup, ReduceOp + +from sglang.srt import _custom_ops as ops +from sglang.srt.utils import is_cuda, is_hip + +logger = logging.getLogger(__name__) + +_is_cuda = is_cuda() +_is_hip = is_hip() + +mscclpp_is_available = False +if _is_hip: + # TODO(zyksir): mscclpp is untested on AMD and therefore disabled. + mscclpp_is_available = False +if _is_cuda: + try: + import sgl_kernel + + mscclpp_is_available = True + except: + mscclpp_is_available = False + + +class MscclContextSelection(IntEnum): + MSCCL1SHOT1NODELL = 1 + MSCCL1SHOT2NODELL = 2 + + +def mscclpp_is_weak_contiguous(inp: torch.Tensor): + return inp.is_contiguous() or ( + inp.storage().nbytes() - inp.storage_offset() * inp.element_size() + == inp.numel() * inp.element_size() + ) + + +def mscclpp_convert_to_bytes(size_str): + """ + Converts a human-readable size string (e.g., "1MB", "2.5kb", "3 GB") + into the equivalent number of bytes using binary units. + + Args: + size_str (str): A string representing size with unit (KB, MB, GB). + + Returns: + int: Number of bytes. + """ + size_str = size_str.strip().lower() + + if not size_str: + raise ValueError("Empty input string") + + # Extract numeric part and unit + for i in range(len(size_str)): + if not size_str[i].isdigit() and size_str[i] != ".": + break + num_str = size_str[:i] + unit = size_str[i:].strip() + + try: + num = float(num_str) + except ValueError: + raise ValueError(f"Invalid numeric value in '{size_str}'") + + # Conversion factors + if unit == "b": + return int(num) + elif unit == "kb": + return int(num * 1024) + elif unit == "mb": + return int(num * 1024 * 1024) + elif unit == "gb": + return int(num * 1024 * 1024 * 1024) + else: + raise ValueError(f"Unsupported unit: {unit}, support B, KB, MB, GB only") + + +def mscclpp_bench_time(func, test_niter: int = 10, warmup_niter: int = 2): + # warmup + for _ in range(warmup_niter): + func() + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + dist.barrier() + start_event.record() + for _ in range(test_niter): + func() + end_event.record() + end_event.synchronize() + func_cost_us = start_event.elapsed_time(end_event) / test_niter * 1000 + return func_cost_us + + +class PyMscclppCommunicator: + _SUPPORTED_WORLD_SIZES = [8, 16] + _MAX_BYTES = mscclpp_convert_to_bytes(os.getenv("SGLANG_MSCCLPP_MAX_BYTES", "1MB")) + _SUPPORTED_DTYPE = [torch.float, torch.float16, torch.bfloat16] + + # max_bytes: max supported mscclpp allreduce size + # in A100 mscclpp is faster than nccl only under condition of msg size smaller than1MB + def __init__( + self, + group: ProcessGroup, + device: Union[int, str, torch.device], + max_bytes=_MAX_BYTES, + ) -> None: + """ + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the CustomAllreduce to. If None, + it will be bind to f"cuda:{local_rank}". + It is the caller's responsibility to make sure each communicator + is bind to a unique device, and all communicators in this group + are in the same node. + """ + self._IS_CAPTURING = False + self.disabled = True + + if not mscclpp_is_available: + # disable because of missing mscclpp library + # e.g. in a non-cuda environment + return + + self.group = group + + assert ( + dist.get_backend(group) != dist.Backend.NCCL + ), "CustomAllreduce should be attached to a non-NCCL group." + + rank = dist.get_rank(group=self.group) + world_size = dist.get_world_size(group=self.group) + if world_size == 1: + # No need to initialize mscclpp for single GPU case. + return + + if world_size not in PyMscclppCommunicator._SUPPORTED_WORLD_SIZES: + logger.warning( + "PyMscclpp is disabled due to an unsupported world" + " size: %d. Supported world sizes: %s. To silence this " + "warning, specify disable_mscclpp=True explicitly.", + world_size, + str(PyMscclppCommunicator._SUPPORTED_WORLD_SIZES), + ) + return + + self.ranks = torch.distributed.get_process_group_ranks(group) + self.nranks_per_node = torch.cuda.device_count() + # for now mscclpp with stride in the communicator is not tested + if not (abs(self.ranks[-1] - self.ranks[0]) == world_size - 1): + logger.warning( + "PyMscclpp is disabled due to an unsupported group %s." + "Please ensure all ranks in the group are consecutive." + "To silence this warning, specify disable_mscclpp=True explicitly.", + str(self.ranks), + ) + return + + if isinstance(device, int): + device = torch.device(f"cuda:{device}") + elif isinstance(device, str): + device = torch.device(device) + # now `device` is a `torch.device` object + assert isinstance(device, torch.device) + self.device = device + + self.max_bytes = max_bytes + self.rank = rank + self.world_size = world_size + + if dist.get_rank(group) == 0: + unique_id = [ops.mscclpp_generate_unique_id()] + else: + unique_id = [None] + dist.broadcast_object_list(unique_id, src=self.ranks[0], group=self.group) + self.unique_id = unique_id[0] + self.rank_to_node, self.rank_to_ib = list(range(world_size)), list( + range(world_size) + ) + for r in range(world_size): + self.rank_to_node[r] = r // 8 + self.rank_to_ib[r] = self.rank % 8 + + self._context = None + self.context_selection = None + self.msg_size_for_finetune = [ + 2**i for i in range(10, math.floor(math.log2(self.max_bytes)) + 1) + ] + self.msg_size2best_config = {} + if world_size == 8: + self.context_selection = MscclContextSelection.MSCCL1SHOT1NODELL + elif world_size == 16: + self.context_selection = MscclContextSelection.MSCCL1SHOT2NODELL + if not _is_hip: + self.scratch = torch.empty( + self.max_bytes * 8, + dtype=torch.uint8, + device=self.device, + ) + self.put_buffer = torch.empty( + self.max_bytes * 8 // self.nranks_per_node, + dtype=torch.uint8, + device=self.device, + ) + self._context = ops.mscclpp_init_context( + self.unique_id, + self.rank, + self.world_size, + self.scratch, + self.put_buffer, + self.nranks_per_node, + self.rank_to_node, + self.rank_to_ib, + int(self.context_selection), + ) + else: + raise NotImplementedError("HIP Mscclpp is not supported yet.") + + self.msg_size2best_config = {} + self.pre_tune_config() + if dist.get_rank(group) == 0: + msg_size2best_config = [self.msg_size2best_config] + else: + msg_size2best_config = [None] + dist.broadcast_object_list( + msg_size2best_config, src=self.ranks[0], group=self.group + ) + self.msg_size2best_config = msg_size2best_config[0] + + # PyMscclpp is enabled only in cuda graph + self.disabled = True + + def pre_tune_config(self, dtype=torch.bfloat16) -> bool: + logger.debug(f"start to pre-tune configs for rank {self.rank}") + nthreads_to_try = [256, 512, 1024] + nblocks_to_try = [21, 42, 84] + inp_randn = torch.ones( + self.msg_size_for_finetune[-1] // dtype.itemsize, dtype=dtype, device="cuda" + ) + oup_randn = torch.empty_like(inp_randn) + for msg_size in self.msg_size_for_finetune: + mock_inp, mock_outp = ( + inp_randn[: msg_size // dtype.itemsize], + oup_randn[: msg_size // dtype.itemsize], + ) + best_config, best_time = None, None + for nthreads in nthreads_to_try: + for nblocks in nblocks_to_try: + cur_cost = mscclpp_bench_time( + lambda: ops.mscclpp_allreduce( + self._context, mock_inp, mock_outp, nthreads, nblocks + ) + ) + if best_time is None or cur_cost < best_time: + best_config = (nthreads, nblocks) + best_time = cur_cost + self.msg_size2best_config[msg_size] = best_config + if self.rank == 0: + logger.debug( + f"for msg_size {msg_size}, best_config: {best_config}, best_time: {best_time}us" + ) + + def should_mscclpp_allreduce( + self, inp: torch.Tensor, op: ReduceOp = ReduceOp.SUM + ) -> bool: + if self.disabled or self._context is None: + return False + if inp.dtype not in PyMscclppCommunicator._SUPPORTED_DTYPE: + return False + if not mscclpp_is_weak_contiguous(inp): + return False + # only support sum op + if op != ReduceOp.SUM: + return False + if inp.numel() * inp.element_size() > self.max_bytes: + return False + return True + + def all_reduce(self, tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM): + if self._IS_CAPTURING: + if torch.cuda.is_current_stream_capturing(): + self.graph_input_set.add((tensor.dtype, tensor.numel())) + msg_size = tensor.numel() * tensor.itemsize + index = bisect.bisect_left(self.msg_size_for_finetune, msg_size) + msg_size_finetune = self.msg_size_for_finetune[index] + nthreads, nblocks = self.msg_size2best_config[msg_size_finetune] + result = torch.empty_like(tensor) + ops.mscclpp_allreduce(self._context, tensor, result, nthreads, nblocks) + return result + + @contextmanager + def change_state( + self, + enable: Optional[bool] = None, + ): + if enable is None: + # guess a default value when not specified + enable = self.available + + old_disable = self.disabled + self.disabled = not enable + + yield + + self.disabled = old_disable diff --git a/python/sglang/srt/distributed/device_communicators/pynccl.py b/python/sglang/srt/distributed/device_communicators/pynccl.py new file mode 100644 index 000000000..fbb59c477 --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/pynccl.py @@ -0,0 +1,341 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/pynccl.py + +import logging +from contextlib import contextmanager +from typing import Optional, Union + +# ===================== import region ===================== +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup, ReduceOp + +from sglang.srt.distributed.device_communicators.pynccl_wrapper import ( + NCCLLibrary, + buffer_type, + cudaStream_t, + ncclComm_t, + ncclDataTypeEnum, + ncclRedOpTypeEnum, + ncclUniqueId, +) +from sglang.srt.distributed.utils import StatelessProcessGroup + +logger = logging.getLogger(__name__) + + +class PyNcclCommunicator: + + def __init__( + self, + group: Union[ProcessGroup, StatelessProcessGroup], + device: Union[int, str, torch.device], + library_path: Optional[str] = None, + ): + """ + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the PyNcclCommunicator to. If None, + it will be bind to f"cuda:{local_rank}". + library_path: the path to the NCCL library. If None, it will + use the default library path. + It is the caller's responsibility to make sure each communicator + is bind to a unique device. + """ + if not isinstance(group, StatelessProcessGroup): + assert dist.is_initialized() + assert ( + dist.get_backend(group) != dist.Backend.NCCL + ), "PyNcclCommunicator should be attached to a non-NCCL group." + # note: this rank is the rank in the group + self.rank = dist.get_rank(group) + self.world_size = dist.get_world_size(group) + else: + self.rank = group.rank + self.world_size = group.world_size + + self.group = group + + # if world_size == 1, no need to create communicator + if self.world_size == 1: + self.available = False + self.disabled = True + self.stream = None + return + try: + self.nccl = NCCLLibrary(library_path) + except Exception: + # disable because of missing NCCL library + # e.g. in a non-GPU environment + self.available = False + self.disabled = True + self.stream = None + return + + self.available = True + self.disabled = False + + self.nccl_version = self.nccl.ncclGetRawVersion() + if self.rank == 0: + logger.info("sglang is using nccl==%s", self.nccl.ncclGetVersion()) + + if self.rank == 0: + # get the unique id from NCCL + self.unique_id = self.nccl.ncclGetUniqueId() + else: + # construct an empty unique id + self.unique_id = ncclUniqueId() + + if not isinstance(group, StatelessProcessGroup): + tensor = torch.ByteTensor(list(self.unique_id.internal)) + ranks = dist.get_process_group_ranks(group) + # arg `src` in `broadcast` is the global rank + dist.broadcast(tensor, src=ranks[0], group=group) + byte_list = tensor.tolist() + for i, byte in enumerate(byte_list): + self.unique_id.internal[i] = byte + else: + self.unique_id = group.broadcast_obj(self.unique_id, src=0) + if isinstance(device, int): + device = torch.device(f"cuda:{device}") + elif isinstance(device, str): + device = torch.device(device) + # now `device` is a `torch.device` object + assert isinstance(device, torch.device) + self.device = device + # nccl communicator and stream will use this device + # `torch.cuda.device` is a context manager that changes the + # current cuda device to the specified one + with torch.cuda.device(device): + self.comm: ncclComm_t = self.nccl.ncclCommInitRank( + self.world_size, self.unique_id, self.rank + ) + self.stream = torch.cuda.Stream() + + # A small all_reduce for warmup. + data = torch.zeros(1, device=device) + self.all_reduce(data) + self.stream.synchronize() + del data + + # by default it is disabled, e.g. in profiling models and prefill phase. + # to use it, use under `with obj.change_state(enable=True)`, usually + # when we are using CUDA graph. + self.disabled = True + + def all_reduce( + self, tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None + ): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = self.stream + self.nccl.ncclAllReduce( + buffer_type(tensor.data_ptr()), + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def all_gather( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + stream=None, + sizes: Optional[list[int]] = None, + ): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert input_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {input_tensor.device}" + ) + if stream is None: + stream = self.stream + + if sizes is not None: + split_offset = 0 + + self.nccl.ncclGroupStart() + for root, split_size in enumerate(sizes): + dst_slice = output_tensor[split_offset : split_offset + split_size] + self.nccl.ncclBroadcast( + buffer_type(input_tensor.data_ptr()), + buffer_type(dst_slice.data_ptr()), + dst_slice.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + root, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + split_offset += split_size + self.nccl.ncclGroupEnd() + else: + self.nccl.ncclAllGather( + buffer_type(input_tensor.data_ptr()), + buffer_type(output_tensor.data_ptr()), + input_tensor.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def reduce_scatter( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + op: ReduceOp = ReduceOp.SUM, + stream=None, + sizes: Optional[list[int]] = None, + ): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert input_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {input_tensor.device}" + ) + if stream is None: + stream = self.stream + + if sizes is not None: + split_offset = 0 + self.nccl.ncclGroupStart() + for root, split_size in enumerate(sizes): + chunk = input_tensor[split_offset : split_offset + split_size, ...] + + self.nccl.ncclReduce( + buffer_type(chunk.data_ptr()), + buffer_type(output_tensor.data_ptr()), + chunk.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), + root, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + split_offset += split_size + self.nccl.ncclGroupEnd() + else: + self.nccl.ncclReduceScatter( + buffer_type(input_tensor.data_ptr()), + buffer_type(output_tensor.data_ptr()), + output_tensor.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def send(self, tensor: torch.Tensor, dst: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = self.stream + self.nccl.ncclSend( + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + dst, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def recv(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = self.stream + self.nccl.ncclRecv( + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + src, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def broadcast(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = self.stream + if src == self.rank: + sendbuff = buffer_type(tensor.data_ptr()) + # NCCL requires the sender also to have a receive buffer + recvbuff = buffer_type(tensor.data_ptr()) + else: + sendbuff = buffer_type() + recvbuff = buffer_type(tensor.data_ptr()) + self.nccl.ncclBroadcast( + sendbuff, + recvbuff, + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + src, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def register_comm_window_raw(self, ptr: int, size: int): + return self.nccl.ncclCommWindowRegister(self.comm, buffer_type(ptr), size, 1) + + def deregister_comm_window(self, window): + return self.nccl.ncclCommWindowDeregister(self.comm, window) + + def group_start(self): + self.nccl.ncclGroupStart() + + def group_end(self): + self.nccl.ncclGroupEnd() + + @contextmanager + def change_state( + self, enable: Optional[bool] = None, stream: Optional[torch.cuda.Stream] = None + ): + """ + A context manager to change the state of the communicator. + """ + if enable is None: + # guess a default value when not specified + enable = self.available + + if stream is None: + stream = self.stream + + old_disable = self.disabled + old_stream = self.stream + + self.stream = stream + self.disabled = not enable + yield + + self.disabled = old_disable + self.stream = old_stream diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py new file mode 100644 index 000000000..d7274cf2c --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py @@ -0,0 +1,133 @@ +import tempfile + +import torch +from packaging import version +from torch.cuda.memory import CUDAPluggableAllocator + +from sglang.srt.distributed.parallel_state import GroupCoordinator +from sglang.srt.managers.schedule_batch import global_server_args_dict + +nccl_allocator_source = """ +#include +extern "C" { + +void* nccl_alloc_plug(size_t size, int device, void* stream) { + void* ptr; + ncclResult_t err = ncclMemAlloc(&ptr, size); + return ptr; + +} + +void nccl_free_plug(void* ptr, size_t size, int device, void* stream) { + ncclResult_t err = ncclMemFree(ptr); +} + +} +""" + +_allocator = None +_mem_pool = None +_registered_base_addrs = set() +_graph_pool_id = None + + +def is_symmetric_memory_enabled(): + return global_server_args_dict["enable_symm_mem"] + + +def set_graph_pool_id(graph_pool_id): + global _graph_pool_id + _graph_pool_id = graph_pool_id + + +def get_nccl_mem_pool(): + global _allocator, _mem_pool + if _mem_pool is None: + out_dir = tempfile.gettempdir() + nccl_allocator_libname = "nccl_allocator" + torch.utils.cpp_extension.load_inline( + name=nccl_allocator_libname, + cpp_sources=nccl_allocator_source, + with_cuda=True, + extra_ldflags=["-lnccl"], + verbose=True, + is_python_module=False, + build_directory=out_dir, + ) + _allocator = CUDAPluggableAllocator( + f"{out_dir}/{nccl_allocator_libname}.so", + "nccl_alloc_plug", + "nccl_free_plug", + ).allocator() + _mem_pool = torch.cuda.MemPool(_allocator) + return _mem_pool + + +class use_symmetric_memory: + def __init__(self, group_coordinator: GroupCoordinator): + if not is_symmetric_memory_enabled(): + self.group_coordinator = None + self._mem_pool_ctx = None + self.is_graph_capture = None + self.device = None + self.pre_2_8_0 = None + else: + self.group_coordinator = group_coordinator + self._mem_pool_ctx = torch.cuda.use_mem_pool(get_nccl_mem_pool()) + self.is_graph_capture = torch.cuda.is_current_stream_capturing() + self.device = torch.cuda.current_device() + self.pre_2_8_0 = version.parse(torch.__version__) < version.parse("2.8.0") + + def __enter__(self): + if not is_symmetric_memory_enabled(): + return self + assert ( + self.group_coordinator.pynccl_comm is not None + ), f"Symmetric memory requires pynccl to be enabled in group '{self.group_coordinator.group_name}'" + assert ( + self.group_coordinator.pynccl_comm.nccl_version >= 22703 + ), "NCCL version 2.27.3 or higher is required for NCCL symmetric memory" + if self.is_graph_capture: + assert ( + _graph_pool_id is not None + ), "graph_pool_id is not set under graph capture" + # Pause graph memory pool to use symmetric memory with cuda graph + if self.pre_2_8_0: + torch._C._cuda_endAllocateCurrentStreamToPool( + self.device, _graph_pool_id + ) + else: + torch._C._cuda_endAllocateToPool(self.device, _graph_pool_id) + self._mem_pool_ctx.__enter__() + return self + + def tag(self, tensor: torch.Tensor): + if not is_symmetric_memory_enabled(): + return + tensor.symmetric_memory = True + + def __exit__(self, exc_type, exc_val, exc_tb): + if not is_symmetric_memory_enabled(): + return + global _registered_base_addrs + self._mem_pool_ctx.__exit__(exc_type, exc_val, exc_tb) + for segment in get_nccl_mem_pool().snapshot(): + if segment["address"] not in _registered_base_addrs: + if segment["stream"] == 0 and self.pre_2_8_0: + # PyTorch version < 2.8.0 has a multi-thread MemPool bug + # See https://github.com/pytorch/pytorch/issues/152861 + # Fixed at https://github.com/pytorch/pytorch/commit/f01e628e3b31852983ab30b25bf251f557ba9c0b + # WAR is to skip allocations on the default stream since the forward_pass thread always runs on a custom stream + continue + self.group_coordinator.pynccl_comm.register_comm_window_raw( + segment["address"], segment["total_size"] + ) + _registered_base_addrs.add(segment["address"]) + + if self.is_graph_capture: + if self.pre_2_8_0: + torch._C._cuda_beginAllocateToPool(self.device, _graph_pool_id) + else: + torch._C._cuda_beginAllocateCurrentThreadToPool( + self.device, _graph_pool_id + ) diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py b/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py new file mode 100644 index 000000000..579811777 --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py @@ -0,0 +1,563 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/pynccl.py + +# This file is a pure Python wrapper for the NCCL library. +# The main purpose is to use NCCL combined with CUDA graph. +# Before writing this script, we tried the following approach: +# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself +# often gets stuck when initializing the NCCL communicator. +# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce` +# contains many other potential cuda APIs, that are not allowed during +# capturing the CUDA graph. For further details, please check +# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ . +# +# Another rejected idea is to write a C/C++ binding for NCCL. It is usually +# doable, but we often encounter issues related with nccl versions, and need +# to switch between different versions of NCCL. See +# https://github.com/NVIDIA/nccl/issues/1234 for more details. +# A C/C++ binding is not flexible enough to handle this. It requires +# recompilation of the code every time we want to switch between different +# versions. This current implementation, with a **pure** Python wrapper, is +# more flexible. We can easily switch between different versions of NCCL by +# changing the environment variable `SGLANG_NCCL_SO_PATH`, or the `so_file` +# variable in the code. + +import ctypes +import logging +import os +import platform +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +import torch +from torch.distributed import ReduceOp + +logger = logging.getLogger(__name__) + + +def find_nccl_library() -> str: + """ + We either use the library file specified by the `SGLANG_NCCL_SO_PATH` + environment variable, or we find the library file brought by PyTorch. + After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be + found by `ctypes` automatically. + """ + + # so_file can be set to None in sglang + so_file = os.environ.get("SGLANG_NCCL_SO_PATH", None) + + # manually load the nccl library + if so_file: + logger.info( + "Found nccl from environment variable SGLANG_NCCL_SO_PATH=%s", so_file + ) + else: + if torch.version.cuda is not None: + so_file = "libnccl.so.2" + elif torch.version.hip is not None: + so_file = "librccl.so.1" + else: + raise ValueError("NCCL only supports CUDA and ROCm backends.") + logger.debug("Found nccl from library %s", so_file) + return so_file + + +# === export types and functions from nccl to Python === +# for the original nccl definition, please check +# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in + +ncclResult_t = ctypes.c_int +ncclComm_t = ctypes.c_void_p +ncclWindow_t = ctypes.c_void_p + + +class ncclUniqueId(ctypes.Structure): + _fields_ = [("internal", ctypes.c_byte * 128)] + + +cudaStream_t = ctypes.c_void_p +buffer_type = ctypes.c_void_p + +ncclDataType_t = ctypes.c_int + + +class ncclDataTypeEnum: + ncclInt8 = 0 + ncclChar = 0 + ncclUint8 = 1 + ncclInt32 = 2 + ncclInt = 2 + ncclUint32 = 3 + ncclInt64 = 4 + ncclUint64 = 5 + ncclFloat16 = 6 + ncclHalf = 6 + ncclFloat32 = 7 + ncclFloat = 7 + ncclFloat64 = 8 + ncclDouble = 8 + ncclBfloat16 = 9 + ncclNumTypes = 10 + + @classmethod + def from_torch(cls, dtype: torch.dtype) -> int: + if dtype == torch.int8: + return cls.ncclInt8 + if dtype == torch.uint8: + return cls.ncclUint8 + if dtype == torch.int32: + return cls.ncclInt32 + if dtype == torch.int64: + return cls.ncclInt64 + if dtype == torch.float16: + return cls.ncclFloat16 + if dtype == torch.float32: + return cls.ncclFloat32 + if dtype == torch.float64: + return cls.ncclFloat64 + if dtype == torch.bfloat16: + return cls.ncclBfloat16 + raise ValueError(f"Unsupported dtype: {dtype}") + + +ncclRedOp_t = ctypes.c_int + + +class ncclRedOpTypeEnum: + ncclSum = 0 + ncclProd = 1 + ncclMax = 2 + ncclMin = 3 + ncclAvg = 4 + ncclNumOps = 5 + + @classmethod + def from_torch(cls, op: ReduceOp) -> int: + if op == ReduceOp.SUM: + return cls.ncclSum + if op == ReduceOp.PRODUCT: + return cls.ncclProd + if op == ReduceOp.MAX: + return cls.ncclMax + if op == ReduceOp.MIN: + return cls.ncclMin + if op == ReduceOp.AVG: + return cls.ncclAvg + raise ValueError(f"Unsupported op: {op}") + + +@dataclass +class Function: + name: str + restype: Any + argtypes: List[Any] + + +class NCCLLibrary: + exported_functions = [ + # const char* ncclGetErrorString(ncclResult_t result) + Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]), + # ncclResult_t ncclGetVersion(int *version); + Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]), + # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId); + Function("ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)]), + # ncclResult_t ncclCommInitRank( + # ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank); + # note that ncclComm_t is a pointer type, so the first argument + # is a pointer to a pointer + Function( + "ncclCommInitRank", + ncclResult_t, + [ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int], + ), + # ncclResult_t ncclAllReduce( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclAllReduce", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclRedOp_t, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclAllGather( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclAllGather", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclReduce( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, int root, + # ncclComm_t comm, cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclReduce", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclRedOp_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclReduceScatter( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclReduceScatter", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclRedOp_t, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclSend( + # const void* sendbuff, size_t count, ncclDataType_t datatype, + # int dest, ncclComm_t comm, cudaStream_t stream); + Function( + "ncclSend", + ncclResult_t, + [ + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclRecv( + # void* recvbuff, size_t count, ncclDataType_t datatype, + # int src, ncclComm_t comm, cudaStream_t stream); + Function( + "ncclRecv", + ncclResult_t, + [ + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclBroadcast( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, int root, ncclComm_t comm, + # cudaStream_t stream); + Function( + "ncclBroadcast", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # be cautious! this is a collective call, it will block until all + # processes in the communicator have called this function. + # because Python object destruction can happen in random order, + # it is better not to call it at all. + # ncclResult_t ncclCommDestroy(ncclComm_t comm); + Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]), + # ncclResult_t ncclGroupStart(); + Function("ncclGroupStart", ncclResult_t, []), + # ncclResult_t ncclGroupEnd(); + Function("ncclGroupEnd", ncclResult_t, []), + ] + + exported_functions_symm_mem = [ + # ncclResult_t ncclCommWindowRegister(ncclComm_t comm, void* buff, size_t size, ncclWindow_t* win, int winFlags); + Function( + "ncclCommWindowRegister", + ncclResult_t, + [ + ncclComm_t, + buffer_type, + ctypes.c_size_t, + ctypes.POINTER(ncclWindow_t), + ctypes.c_int, + ], + ), + # ncclResult_t ncclCommWindowDeregister(ncclComm_t comm, ncclWindow_t win); + Function("ncclCommWindowDeregister", ncclResult_t, [ncclComm_t, ncclWindow_t]), + ] + + # class attribute to store the mapping from the path to the library + # to avoid loading the same library multiple times + path_to_library_cache: Dict[str, Any] = {} + + # class attribute to store the mapping from library path + # to the corresponding dictionary + path_to_dict_mapping: Dict[str, Dict[str, Any]] = {} + + def __init__(self, so_file: Optional[str] = None): + + so_file = so_file or find_nccl_library() + + try: + if so_file not in NCCLLibrary.path_to_dict_mapping: + lib = ctypes.CDLL(so_file) + NCCLLibrary.path_to_library_cache[so_file] = lib + self.lib = NCCLLibrary.path_to_library_cache[so_file] + except Exception as e: + logger.error( + "Failed to load NCCL library from %s ." + "It is expected if you are not running on NVIDIA/AMD GPUs." + "Otherwise, the nccl library might not exist, be corrupted " + "or it does not support the current platform %s." + "If you already have the library, please set the " + "environment variable SGLANG_NCCL_SO_PATH" + " to point to the correct nccl library path.", + so_file, + platform.platform(), + ) + raise e + + if so_file not in NCCLLibrary.path_to_dict_mapping: + _funcs: Dict[str, Any] = {} + exported_functions = NCCLLibrary.exported_functions + if hasattr(self.lib, "ncclCommWindowRegister"): + exported_functions.extend(NCCLLibrary.exported_functions_symm_mem) + for func in exported_functions: + f = getattr(self.lib, func.name) + f.restype = func.restype + f.argtypes = func.argtypes + _funcs[func.name] = f + NCCLLibrary.path_to_dict_mapping[so_file] = _funcs + self._funcs = NCCLLibrary.path_to_dict_mapping[so_file] + + def ncclGetErrorString(self, result: ncclResult_t) -> str: + return self._funcs["ncclGetErrorString"](result).decode("utf-8") + + def NCCL_CHECK(self, result: ncclResult_t) -> None: + if result != 0: + error_str = self.ncclGetErrorString(result) + raise RuntimeError(f"NCCL error: {error_str}") + + def ncclGetRawVersion(self) -> int: + version = ctypes.c_int() + self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version))) + # something like 21903 + return version.value + + def ncclGetVersion(self) -> str: + version_str = str(self.ncclGetRawVersion()) + # something like 21903 --> "2.19.3" + major = version_str[0].lstrip("0") + minor = version_str[1:3].lstrip("0") + patch = version_str[3:].lstrip("0") + return f"{major}.{minor}.{patch}" + + def ncclGetUniqueId(self) -> ncclUniqueId: + unique_id = ncclUniqueId() + self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id))) + return unique_id + + def ncclCommInitRank( + self, world_size: int, unique_id: ncclUniqueId, rank: int + ) -> ncclComm_t: + comm = ncclComm_t() + self.NCCL_CHECK( + self._funcs["ncclCommInitRank"]( + ctypes.byref(comm), world_size, unique_id, rank + ) + ) + return comm + + def ncclAllReduce( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + op: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclAllReduce"]( + sendbuff, recvbuff, count, datatype, op, comm, stream + ) + ) + + def ncclReduce( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + op: int, + root: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclReduce"]( + sendbuff, recvbuff, count, datatype, op, root, comm, stream + ) + ) + + def ncclReduceScatter( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + op: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclReduceScatter"]( + sendbuff, recvbuff, count, datatype, op, comm, stream + ) + ) + + def ncclAllGather( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # which is an aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclAllGather"]( + sendbuff, recvbuff, count, datatype, comm, stream + ) + ) + + def ncclSend( + self, + sendbuff: buffer_type, + count: int, + datatype: int, + dest: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + self.NCCL_CHECK( + self._funcs["ncclSend"](sendbuff, count, datatype, dest, comm, stream) + ) + + def ncclRecv( + self, + recvbuff: buffer_type, + count: int, + datatype: int, + src: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + self.NCCL_CHECK( + self._funcs["ncclRecv"](recvbuff, count, datatype, src, comm, stream) + ) + + def ncclBroadcast( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + root: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + self.NCCL_CHECK( + self._funcs["ncclBroadcast"]( + sendbuff, recvbuff, count, datatype, root, comm, stream + ) + ) + + def ncclCommDestroy(self, comm: ncclComm_t) -> None: + self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm)) + + def ncclCommWindowRegister( + self, comm: ncclComm_t, buff: buffer_type, size: int, win_flags: int + ) -> ncclWindow_t: + window = ncclWindow_t() + self.NCCL_CHECK( + self._funcs["ncclCommWindowRegister"]( + comm, buff, size, ctypes.byref(window), win_flags + ) + ) + return window + + def ncclCommWindowDeregister(self, comm: ncclComm_t, window: ncclWindow_t) -> None: + self.NCCL_CHECK(self._funcs["ncclCommWindowDeregister"](comm, window)) + + def ncclGroupStart(self) -> None: + self.NCCL_CHECK(self._funcs["ncclGroupStart"]()) + + def ncclGroupEnd(self) -> None: + self.NCCL_CHECK(self._funcs["ncclGroupEnd"]()) + + +__all__ = [ + "NCCLLibrary", + "ncclDataTypeEnum", + "ncclRedOpTypeEnum", + "ncclUniqueId", + "ncclComm_t", + "cudaStream_t", + "buffer_type", +] diff --git a/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py b/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py new file mode 100644 index 000000000..0113c432d --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 + +import logging +import os +from enum import Enum +from typing import Union + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from sglang.srt import _custom_ops as ops +from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import ( + is_full_nvlink, + is_weak_contiguous, +) +from sglang.srt.distributed.parallel_state import in_the_same_node_as +from sglang.srt.utils import is_cuda, is_hip + +logger = logging.getLogger(__name__) + +_is_cuda = is_cuda() +_is_hip = is_hip() + + +try: + ops.qr_max_size() + quick_ar = True +except Exception: + # For CPUs and CUDA + quick_ar = False + + +def qr_rocm_arch_available(): + if not _is_hip: + return False + try: + props = torch.cuda.get_device_properties(0) + gcn_arch = getattr(props, "gcnArchName", "") + supported_archs = ["gfx94", "gfx95"] + return any(gfx in gcn_arch for gfx in supported_archs) + except Exception as e: + logger.warning("Failed to determine ROCm for quick allreduce: %s", e) + return False + + +class QuickReduceRegime(Enum): + FP = 0 + INT8 = 1 + INT6 = 2 + INT4 = 3 + NONE = 4 + + +MB = 1024 * 1024 + + +class QuickAllReduce: + + _SUPPORTED_WORLD_SIZES = [2, 4, 8] + _SUPPORTED_DTYPES = [torch.float16, torch.bfloat16] + # The following data is based on kernel tests. + # In this order [FP, INT8, INT6, INT4]. + _QR_MIN_SIZE = { + (torch.float16, 2): [1 * MB, 2 * MB, 2 * MB, 1 * MB], + (torch.float16, 4): [1 * MB, 16 * MB, 4 * MB, 2 * MB], + (torch.float16, 8): [16 * MB, 4 * MB, 4 * MB, 2 * MB], + (torch.bfloat16, 2): [2 * MB, 8 * MB, 8 * MB, 8 * MB], + (torch.bfloat16, 4): [8 * MB, 64 * MB, 64 * MB, 16 * MB], + (torch.bfloat16, 8): [16 * MB, 2048 * MB, 2048 * MB, 2048 * MB], + } + + def __init__( + self, group: ProcessGroup, device: Union[int, str, torch.device] + ) -> None: + """ + Custom allreduce provides non-destructive acceleration and is + available for CUDA and ROCm MI300 series. + Custom quick allreduce leverages quantization for further + acceleration on ROCm. It currently supports Q8, Q6, and Q4 + quantization formats and FP(float16, bfloat16). + Quick allreduce is designed as a complement to custom allreduce. + Its initialization requires even stricter conditions. + Only the ROCm MI300 series is supported for quick allreduce at + this time. + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the CustomAllreduce to. If None, + it will be bind to f"cuda:{local_rank}". + It is the caller's responsibility to make sure each communicator + is bind to a unique device, and all communicators in this group + are in the same node. + """ + self.disabled = True + if not qr_rocm_arch_available(): + logger.debug( + "Custom quick allreduce is only supported on ROCm MI300 series." + ) + return + + if not quick_ar: + # disable because of missing quick reduce library + # e.g. in a cuda environment + logger.info( + "Custom quick allreduce is disabled because " + "of missing custom quick allreduce library" + ) + return + + self.group = group + assert ( + dist.get_backend(group) != dist.Backend.NCCL + ), "Custom quick allreduce should be attached to a non-NCCL group." + if not all(in_the_same_node_as(group, source_rank=0)): + # No need to initialize custom quick allreduce for + # multi-node case. + logger.warning( + "Custom quick allreduce is disabled because this " + "process group spans across nodes." + ) + return + rank = dist.get_rank(group=self.group) + world_size = dist.get_world_size(group=self.group) + self.rank = rank + self.world_size = world_size + if world_size == 1: + # No need to initialize QuickReduce for single GPU case. + return + + if world_size not in QuickAllReduce._SUPPORTED_WORLD_SIZES: + logger.warning( + "Custom quick allreduce is disabled due to an " + "unsupported world size: %d. Supported world sizes: %s.", + world_size, + str(QuickAllReduce._SUPPORTED_WORLD_SIZES), + ) + return + + if isinstance(device, int): + device = torch.device(f"cuda:{device}") + elif isinstance(device, str): + device = torch.device(device) + assert isinstance(device, torch.device) + self.device = device + + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None) + if cuda_visible_devices: + device_ids = list(map(int, cuda_visible_devices.split(","))) + else: + device_ids = list(range(torch.cuda.device_count())) + physical_device_id = device_ids[device.index] + tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") + gather_list = [ + torch.tensor([0], dtype=torch.int, device="cpu") + for _ in range(self.world_size) + ] + dist.all_gather(gather_list, tensor, group=self.group) + physical_device_ids = [t.item() for t in gather_list] + + # test nvlink first, this will filter out most of the cases + # where custom quick allreduce is not supported + # this checks hardware and driver support for NVLink + if _is_cuda or _is_hip: + self.fully_connected = is_full_nvlink(physical_device_ids, self.world_size) + if self.world_size > 2 and not self.fully_connected: + logger.debug( + "Custom quick allreduce is disabled because it's not supported " + "on more than two PCIe-only GPUs. " + ) + return + + self.init_quick_all_reduce() + + def init_quick_all_reduce(self): + # On RocM, bfloat16 kernels are slower than fp16 + # due to slower match operations + # If environment variable is set to 1, we convert input to fp16 + self.use_fp16_kernels = int( + os.environ.get("ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", 1) + ) + regime_str = os.environ.get("ROCM_QUICK_REDUCE_QUANTIZATION", "NONE") + if regime_str not in QuickReduceRegime.__members__: + logger.warning( + "Custom quick allreduce:", + f"Invalid quantization level: {regime_str}. " + "Supported levels: " + f"{list(QuickReduceRegime.__members__.keys())}", + ) + return + + if regime_str == "NONE": + logger.debug( + "Custom quick allreduce is disabled based " + "on env variable " + "ROCM_QUICK_REDUCE_QUANTIZATION='NONE'" + ) + return + self.qr_quant_level = QuickReduceRegime[regime_str] + + # TODO: If the dtype is not bfloat16 or then float16, + # quickallreduce should not be created. + + # ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is specified in MB + qr_max_size = int(os.environ.get("ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", 0)) + if qr_max_size > 0: + if qr_max_size < 1: + logger.info( + "You should not set a max_size smaller than 1MB, which can " + "lead to error or degradation to custom allreduce or rccl." + ) + qr_max_size = qr_max_size * MB + # If qr_max_size is None, then 2GB is used by default. + self._ptr = ops.init_custom_qr(self.rank, self.world_size, qr_max_size) + self.qr_max_size = qr_max_size if qr_max_size > 0 else ops.qr_max_size() + self.create_shared_buffer() + self.disabled = False + + def create_shared_buffer(self): + """ + Creates a shared buffer for quickreduce. + Has to be called after init_custom_qr + """ + handle = ops.qr_get_handle(self._ptr) + world_size = dist.get_world_size(group=self.group) + handles = [None] * world_size + dist.all_gather_object(handles, handle, group=self.group) + ops.qr_open_handles(self._ptr, handles) + + def should_quick_allreduce(self, inp: torch.Tensor): + """ + Check if quickreduce is available + """ + if self.disabled: + return False + if inp.dtype not in self._SUPPORTED_DTYPES: + return False + inp_size = inp.numel() * inp.element_size() + # custom quick allreduce requires input byte size to be + # multiples of 16 + if inp_size % 16 != 0: + return False + if not is_weak_contiguous(inp): + return False + dtype = inp.dtype + if self.use_fp16_kernels: + dtype = torch.float16 + return ( + inp_size <= self.qr_max_size + and inp_size + >= self._QR_MIN_SIZE[(dtype, self.world_size)][self.qr_quant_level.value] + ) + + def quick_all_reduce(self, inp: torch.Tensor, *, out: torch.Tensor = None): + """Performs an out-of-place custom quick all reduce.""" + # quick allreduce doesn't require a separate graph mode, + # as QR uses static IPC buffer. + if out is None: + out = torch.empty_like(inp) + ops.qr_all_reduce( + self._ptr, inp, out, self.qr_quant_level.value, self.use_fp16_kernels + ) + return out + + def close(self): + if not self.disabled and getattr(self, "_ptr", None): + if ops is not None: + ops.qr_destroy(self._ptr) + self._ptr = 0 + self.disabled = True + + def __del__(self): + self.close() diff --git a/python/sglang/srt/distributed/device_communicators/shm_broadcast.py b/python/sglang/srt/distributed/device_communicators/shm_broadcast.py new file mode 100644 index 000000000..e5b59e7cc --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/shm_broadcast.py @@ -0,0 +1,509 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/shm_broadcast.py + +import logging +import os +import pickle +import time +from contextlib import contextmanager +from dataclasses import dataclass, field +from multiprocessing import shared_memory +from typing import List, Optional +from unittest.mock import patch + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup +from zmq import IPV6 # type: ignore +from zmq import SUB, SUBSCRIBE, XPUB, XPUB_VERBOSE, Context # type: ignore + +from sglang.srt.utils import ( + format_tcp_address, + get_ip, + get_open_port, + is_valid_ipv6_address, +) + +# SGLANG_RINGBUFFER_WARNING_INTERVAL can be set to 60 +SGLANG_RINGBUFFER_WARNING_INTERVAL = int( + os.environ.get("SGLANG_RINGBUFFER_WARNING_INTERVAL", "60") +) + +logger = logging.getLogger(__name__) + + +class ShmRingBuffer: + + def __init__( + self, + n_reader: int, + max_chunk_bytes: int, + max_chunks: int, + name: Optional[str] = None, + ): + """ + A shared memory ring buffer implementation for broadcast communication. + Essentially, it is a queue where only one will `enqueue` and multiple + will `dequeue`. The max size of each item, together with the max number + of items that can be stored in the buffer are known in advance. + In this case, we don't need to synchronize the access to + the buffer. + + Buffer memory layout: + data metadata + | | + | (current_idx) | (current_idx) + v v + +-------------------------------+----------------------------------------+ + | chunk0 | chunk1 | ... | chunk | metadata0 | metadata1 | ... | metadata | + +-------------------------------+----------------------------------------+ + | max_chunks x max_chunk_bytes | max_chunks x (1 + n_reader) bytes | + + metadata memory layout: each byte is a flag, the first byte is the written + flag, and the rest are reader flags. The flags are set to 0 by default. + +--------------+--------------+--------------+-----+--------------+ + | written_flag | reader0_flag | reader1_flag | ... | readerN_flag | + +--------------+--------------+--------------+-----+--------------+ + + The state of metadata is as follows: + + (case 1) 0???...???: the block is not written yet, cannot read, can write + (case 2) 1000...000: the block is just written, can read, cannot write + (case 3) 1???...???: the block is written and read by some readers, can read if not read, cannot write + (case 4) 1111...111: the block is written and read by all readers, cannot read, can write + + State transition for readers: + + When a reader finds a block that it can read (case 2 or 3), it can yield the block for caller to read. + Only after the caller finishes reading the block, the reader can mark the block as read. + Readers only mark the block as read (from 0 to 1), the writer marks the block as ready to read (from 1 to 0). + + State transition for writer: + + When the writer writes to a block (case 1 or 4), it first resets the written flag to 0, converting either case + to case 1. Then it can yield the block for caller to write. After the caller finishes writing the block, the writer + can reset the reader flags to 0, and mark the block as written (from 0 to 1). + NOTE: the order is important here, first reset the reader flags (so that we are still in case 1), then mark the block as written. The state transition is atomic. If we do it in the reverse order, it will go through case 3 and then back to case 2, and readers might read the intermediate case 3, which is not correct. + + During creation, `name` is None and the buffer is created. We can pass the + created object to other processes by pickling it. The other processes will + get the name of the shared memory and open it, so that they can access the + same shared memory buffer. + """ # noqa + self.n_reader = n_reader + self.metadata_size = 1 + n_reader + self.max_chunk_bytes = max_chunk_bytes + self.max_chunks = max_chunks + self.total_bytes_of_buffer = ( + self.max_chunk_bytes + self.metadata_size + ) * self.max_chunks + self.data_offset = 0 + self.metadata_offset = self.max_chunk_bytes * self.max_chunks + + if name is None: + # we are creating a buffer + self.is_creator = True + self.shared_memory = shared_memory.SharedMemory( + create=True, size=self.total_bytes_of_buffer + ) + # initialize the metadata section to 0 + with memoryview( + self.shared_memory.buf[self.metadata_offset :] + ) as metadata_buffer: + torch.frombuffer(metadata_buffer, dtype=torch.uint8).fill_(0) + else: + # we are opening an existing buffer + self.is_creator = False + # fix to https://stackoverflow.com/q/62748654/9191338 + # Python incorrectly tracks shared memory even if it is not + # created by the process. The following patch is a workaround. + with patch( + "multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None, + ): + try: + self.shared_memory = shared_memory.SharedMemory(name=name) + assert self.shared_memory.size == self.total_bytes_of_buffer + except FileNotFoundError: + # we might deserialize the object in a different node + # in this case, this object is not used, + # and we should suppress the error + pass + + def __reduce__(self): + return ( + self.__class__, + ( + self.n_reader, + self.max_chunk_bytes, + self.max_chunks, + self.shared_memory.name, + ), + ) + + def __del__(self): + if hasattr(self, "shared_memory"): + self.shared_memory.close() + if self.is_creator: + self.shared_memory.unlink() + + @contextmanager + def get_data(self, current_idx: int): + start = self.data_offset + current_idx * self.max_chunk_bytes + end = start + self.max_chunk_bytes + with memoryview(self.shared_memory.buf[start:end]) as buf: + yield buf + + @contextmanager + def get_metadata(self, current_idx: int): + start = self.metadata_offset + current_idx * self.metadata_size + end = start + self.metadata_size + with memoryview(self.shared_memory.buf[start:end]) as buf: + yield buf + + +@dataclass +class Handle: + connect_ip: str + local_reader_ranks: List[int] = field(default_factory=list) + + buffer: Optional[ShmRingBuffer] = None + local_subscribe_port: Optional[int] = None + remote_subscribe_port: Optional[int] = None + + +class MessageQueue: + + def __init__( + self, + n_reader, # number of all readers + n_local_reader, # number of local readers through shared memory + local_reader_ranks: Optional[List[int]] = None, + max_chunk_bytes: int = 1024 * 1024 * 10, + max_chunks: int = 10, + connect_ip: Optional[str] = None, + ): + if local_reader_ranks is None: + local_reader_ranks = list(range(n_local_reader)) + else: + assert len(local_reader_ranks) == n_local_reader + self.n_local_reader = n_local_reader + n_remote_reader = n_reader - n_local_reader + self.n_remote_reader = n_remote_reader + + if connect_ip is None: + connect_ip = get_ip() if n_remote_reader > 0 else "127.0.0.1" + + context = Context() + + if n_local_reader > 0: + # for local readers, we will: + # 1. create a shared memory ring buffer to communicate small data + # 2. create a publish-subscribe socket to communicate large data + self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks) + + # XPUB is very similar to PUB, + # except that it can receive subscription messages + # to confirm the number of subscribers + self.local_socket = context.socket(XPUB) + # set the verbose option so that we can receive every subscription + # message. otherwise, we will only receive the first subscription + # see http://api.zeromq.org/3-3:zmq-setsockopt for more details + self.local_socket.setsockopt(XPUB_VERBOSE, True) + local_subscribe_port = get_open_port() + socket_addr = f"tcp://127.0.0.1:{local_subscribe_port}" + logger.debug("Binding to %s", socket_addr) + self.local_socket.bind(socket_addr) + + self.current_idx = 0 + + else: + self.buffer = None # type: ignore + local_subscribe_port = None + self.local_socket = None + self.current_idx = -1 + + if n_remote_reader > 0: + # for remote readers, we will: + # create a publish-subscribe socket to communicate large data + self.remote_socket = context.socket(XPUB) + self.remote_socket.setsockopt(XPUB_VERBOSE, True) + remote_subscribe_port = get_open_port() + if is_valid_ipv6_address(connect_ip): + self.remote_socket.setsockopt(IPV6, 1) + self.remote_socket.bind( + format_tcp_address(connect_ip, remote_subscribe_port) + ) + + else: + remote_subscribe_port = None + self.remote_socket = None + + self._is_writer = True + self._is_local_reader = False + self.local_reader_rank = -1 + # rank does not matter for remote readers + self._is_remote_reader = False + + self.handle = Handle( + connect_ip=connect_ip, + local_reader_ranks=local_reader_ranks, + buffer=self.buffer, + local_subscribe_port=local_subscribe_port, + remote_subscribe_port=remote_subscribe_port, + ) + + logger.debug("Message queue communication handle: %s", self.handle) + + def export_handle(self) -> Handle: + return self.handle + + @staticmethod + def create_from_handle(handle: Handle, rank) -> "MessageQueue": + self = MessageQueue.__new__(MessageQueue) + self.handle = handle + self._is_writer = False + + context = Context() + + if rank in handle.local_reader_ranks: + assert handle.buffer is not None + self.buffer = handle.buffer + self.current_idx = 0 + self.local_reader_rank = handle.local_reader_ranks.index(rank) + self._is_local_reader = True + self._is_remote_reader = False + + self.local_socket = context.socket(SUB) + self.local_socket.setsockopt_string(SUBSCRIBE, "") + socket_addr = f"tcp://127.0.0.1:{handle.local_subscribe_port}" + logger.debug("Connecting to %s", socket_addr) + self.local_socket.connect(socket_addr) + + self.remote_socket = None + else: + self.buffer = None # type: ignore + self.current_idx = -1 + self.local_reader_rank = -1 + self._is_local_reader = False + self._is_remote_reader = True + + self.local_socket = None + + self.remote_socket = context.socket(SUB) + self.remote_socket.setsockopt_string(SUBSCRIBE, "") + if is_valid_ipv6_address(handle.connect_ip): + self.remote_socket.setsockopt(IPV6, 1) + socket_addr = format_tcp_address( + handle.connect_ip, handle.remote_subscribe_port + ) + logger.debug("Connecting to %s", socket_addr) + self.remote_socket.connect(socket_addr) + + return self + + def wait_until_ready(self): + """This is a collective operation. All processes (including the + readers and the writer) should call this function. + """ + if self._is_writer: + # wait for all readers to connect + + # local readers + for i in range(self.n_local_reader): + # wait for subscription messages from all local readers + self.local_socket.recv() + if self.n_local_reader > 0: + # send a message to all local readers + # to make sure the publish channel is working + self.local_socket.send(b"READY") + + # remote readers + for i in range(self.n_remote_reader): + # wait for subscription messages from all remote readers + self.remote_socket.recv() + if self.n_remote_reader > 0: + # send a message to all remote readers + # to make sure the publish channel is working + self.remote_socket.send(b"READY") + elif self._is_local_reader: + # wait for the writer to send a message + recv = self.local_socket.recv() + assert recv == b"READY" + elif self._is_remote_reader: + # wait for the writer to send a message + recv = self.remote_socket.recv() + assert recv == b"READY" + + @contextmanager + def acquire_write(self): + assert self._is_writer, "Only writers can acquire write" + start_time = time.monotonic() + n_warning = 1 + while True: + with self.buffer.get_metadata(self.current_idx) as metadata_buffer: + read_count = sum(metadata_buffer[1:]) + written_flag = metadata_buffer[0] + if written_flag and read_count != self.buffer.n_reader: + # this block is written and not read by all readers + # for writers, `self.current_idx` is the next block to write + # if this block is not ready to write, + # we need to wait until it is read by all readers + + # Release the processor to other threads + os.sched_yield() + + # if we wait for a long time, we should warn the user + if ( + time.monotonic() - start_time + > SGLANG_RINGBUFFER_WARNING_INTERVAL * n_warning + ): + logger.warning( + "No available block found in %s second. ", + SGLANG_RINGBUFFER_WARNING_INTERVAL, + ) + n_warning += 1 + + continue + # found a block that is either + # (1) not written + # (2) read by all readers + + # mark the block as not written + metadata_buffer[0] = 0 + # let caller write to the buffer + with self.buffer.get_data(self.current_idx) as buf: + yield buf + + # caller has written to the buffer + # NOTE: order is important here + # first set the read flags to 0 + # then set the written flag to 1 + # otherwise, the readers may think they already read the block + for i in range(1, self.buffer.n_reader + 1): + # set read flag to 0, meaning it is not read yet + metadata_buffer[i] = 0 + # mark the block as written + metadata_buffer[0] = 1 + self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks + break + + @contextmanager + def acquire_read(self): + assert self._is_local_reader, "Only readers can acquire read" + start_time = time.monotonic() + n_warning = 1 + while True: + with self.buffer.get_metadata(self.current_idx) as metadata_buffer: + read_flag = metadata_buffer[self.local_reader_rank + 1] + written_flag = metadata_buffer[0] + if not written_flag or read_flag: + # this block is either + # (1) not written + # (2) already read by this reader + + # for readers, `self.current_idx` is the next block to read + # if this block is not ready, + # we need to wait until it is written + + # Release the processor to other threads + os.sched_yield() + + # if we wait for a long time, we should warn the user + if ( + time.monotonic() - start_time + > SGLANG_RINGBUFFER_WARNING_INTERVAL * n_warning + ): + logger.warning( + "No available block found in %s second. ", + SGLANG_RINGBUFFER_WARNING_INTERVAL, + ) + n_warning += 1 + + continue + # found a block that is not read by this reader + # let caller read from the buffer + with self.buffer.get_data(self.current_idx) as buf: + yield buf + + # caller has read from the buffer + # set the read flag + metadata_buffer[self.local_reader_rank + 1] = 1 + self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks + break + + def enqueue(self, obj): + assert self._is_writer, "Only writers can enqueue" + serialized_obj = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) + if self.n_local_reader > 0: + if len(serialized_obj) >= self.buffer.max_chunk_bytes: + with self.acquire_write() as buf: + buf[0] = 1 # overflow + self.local_socket.send(serialized_obj) + else: + with self.acquire_write() as buf: + buf[0] = 0 # not overflow + buf[1 : len(serialized_obj) + 1] = serialized_obj + if self.n_remote_reader > 0: + self.remote_socket.send(serialized_obj) + + def dequeue(self): + if self._is_local_reader: + with self.acquire_read() as buf: + overflow = buf[0] == 1 + if not overflow: + # no need to know the size of serialized object + # pickle format contains the size information internally + # see https://docs.python.org/3/library/pickle.html + obj = pickle.loads(buf[1:]) + if overflow: + recv = self.local_socket.recv() + obj = pickle.loads(recv) + elif self._is_remote_reader: + recv = self.remote_socket.recv() + obj = pickle.loads(recv) + else: + raise RuntimeError("Only readers can dequeue") + return obj + + def broadcast_object(self, obj=None): + if self._is_writer: + self.enqueue(obj) + return obj + else: + return self.dequeue() + + @staticmethod + def create_from_process_group( + pg: ProcessGroup, max_chunk_bytes, max_chunks, writer_rank=0 + ) -> "MessageQueue": + group_rank = dist.get_rank(pg) + group_world_size = dist.get_world_size(pg) + global_ranks = dist.get_process_group_ranks(pg) + + from sglang.srt.distributed.parallel_state import in_the_same_node_as + + status = in_the_same_node_as(pg, source_rank=writer_rank) + same_node_ranks = [i for i, s in enumerate(status) if s] + n_reader = group_world_size - 1 + n_local_reader = len(same_node_ranks) - 1 + local_reader_ranks = [i for i in same_node_ranks if i != writer_rank] + buffer_io: MessageQueue + if group_rank == writer_rank: + buffer_io = MessageQueue( + n_reader=n_reader, + n_local_reader=n_local_reader, + local_reader_ranks=local_reader_ranks, + max_chunk_bytes=max_chunk_bytes, + max_chunks=max_chunks, + ) + handle = buffer_io.export_handle() + dist.broadcast_object_list( + [handle], src=global_ranks[writer_rank], group=pg + ) + else: + recv = [None] + dist.broadcast_object_list(recv, src=global_ranks[writer_rank], group=pg) + handle = recv[0] # type: ignore + buffer_io = MessageQueue.create_from_handle(handle, group_rank) + buffer_io.wait_until_ready() + return buffer_io diff --git a/python/sglang/srt/distributed/device_communicators/xpu_communicator.py b/python/sglang/srt/distributed/device_communicators/xpu_communicator.py new file mode 100644 index 000000000..532279f70 --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/xpu_communicator.py @@ -0,0 +1,48 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/xpu_communicator.py + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from sglang.srt.utils import is_xpu + + +class XpuCommunicator: + + def __init__(self, group: ProcessGroup): + if not is_xpu(): + self.disabled = True + return + self.disabled = False + self.group = group + self.world_size = dist.get_world_size(self.group) + + def all_reduce(self, x: torch.Tensor) -> torch.Tensor: + dist.all_reduce(x, group=self.group) + return x + + def gather( + self, input_: torch.Tensor, rank_in_group: int, dst: int = 0, dim: int = -1 + ): + # For xpu path, gather doesn't work properly together with ray + # cluster so we use all_gather instead for now. + input_size = input_.size() + # Allocate output tensor. + output_tensor = torch.empty( + (self.world_size,) + input_size, dtype=input_.dtype, device=input_.device + ) + # All-gather. + torch.distributed.all_gather_into_tensor( + output_tensor, input_, group=self.group + ) + if rank_in_group == dst: + # Reshape + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + + (self.world_size * input_size[dim],) + + input_size[dim + 1 :] + ) + else: + output_tensor = None + return output_tensor diff --git a/python/sglang/srt/distributed/naive_distributed.py b/python/sglang/srt/distributed/naive_distributed.py new file mode 100644 index 000000000..61165d90c --- /dev/null +++ b/python/sglang/srt/distributed/naive_distributed.py @@ -0,0 +1,112 @@ +import base64 +import os +import pickle +import time +from pathlib import Path +from typing import Any, List, Optional + +import torch + +from sglang.srt.utils import MultiprocessingSerializer + + +class NaiveDistributed: + def __init__(self, rank: int, world_size: int, rendezvous: str): + self._rank = rank + self._world_size = world_size + self._operation_index = 0 + self._directory = Path(rendezvous) + self._directory.mkdir(parents=True, exist_ok=True) + assert 0 <= rank < world_size + + # both barrier to be safe, and as a sanity check + self.barrier() + + def get_rank(self): + return self._rank + + def get_world_size(self): + return self._world_size + + def scatter( + self, tensor: torch.Tensor, scatter_list: List[torch.Tensor], src: int = 0 + ): + if self._rank == src: + assert len(scatter_list) == self._world_size + else: + assert scatter_list is None + + gathered_objects = self.all_gather_object( + dict( + serialized_scatter_list=[ + ( + None + if item_rank == src + else MultiprocessingSerializer.serialize(item) + ) + for item_rank, item in enumerate(scatter_list) + ] + ) + if self._rank == src + else dict() + ) + + remote_serialized_tensor = gathered_objects[src]["serialized_scatter_list"][ + self._rank + ] + if self._rank == src: + assert remote_serialized_tensor is None + remote_tensor = scatter_list[self._rank] + else: + remote_tensor = MultiprocessingSerializer.deserialize( + remote_serialized_tensor + ) + tensor.copy_(remote_tensor) + + # avoid src tensor be deleted too early + self.barrier() + + def all_gather_object(self, obj: Any) -> List[Any]: + self._operation_index += 1 + + text_postfix = "\n" + + def _get_path(interesting_rank: int): + return ( + self._directory + / f"rank{interesting_rank}_op{self._operation_index}.txt" + ) + + _get_path(self._rank).write_text( + base64.b64encode(pickle.dumps(obj)).decode("utf-8") + text_postfix + ) + + def _read_one(interesting_rank: int): + p = _get_path(interesting_rank) + while True: + if p.exists() and (text := p.read_text()).endswith(text_postfix): + return pickle.loads(base64.b64decode(text[: -len(text_postfix)])) + time.sleep(0.001) + + return [ + _read_one(interesting_rank) for interesting_rank in range(self._world_size) + ] + + def barrier(self): + actual_objs = self.all_gather_object(self._rank) + assert actual_objs == list(range(self._world_size)), f"{actual_objs=}" + + +# Can have multi instances if needed +_instance: Optional[NaiveDistributed] = None + + +def get_naive_distributed(): + assert _instance is not None + return _instance + + +def set_naive_distributed(instance: NaiveDistributed): + global _instance + assert _instance is None + _instance = instance diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py new file mode 100644 index 000000000..4f410570d --- /dev/null +++ b/python/sglang/srt/distributed/parallel_state.py @@ -0,0 +1,1776 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/parallel_state.py + +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +"""vLLM distributed state. +It takes over the control of the distributed environment from PyTorch. +The typical workflow is: + +- call `init_distributed_environment` to initialize the distributed environment. +- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to + initialize the model parallel groups. + +- any code dealing with the distributed stuff + +- call `destroy_model_parallel` to destroy the model parallel groups. +- call `destroy_distributed_environment` to destroy the distributed environment. + +If you only need to use the distributed environment without model/pipeline + parallelism, you can skip the model parallel initialization and destruction + steps. +""" +import contextlib +import gc +import logging +import os +import pickle +import weakref +from collections import namedtuple +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from datetime import timedelta +from multiprocessing import shared_memory +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from unittest.mock import patch + +import torch +import torch.distributed +from torch.distributed import Backend, ProcessGroup + +from sglang.srt.utils import ( + direct_register_custom_op, + get_bool_env_var, + get_int_env_var, + is_cpu, + is_cuda_alike, + is_hip, + is_npu, + is_shm_available, + supports_custom_op, +) + +_is_npu = is_npu() +_is_cpu = is_cpu() + +IS_ONE_DEVICE_PER_PROCESS = get_bool_env_var("SGLANG_ONE_DEVICE_PER_PROCESS") + + +@dataclass +class GraphCaptureContext: + stream: torch.cuda.Stream if not _is_npu else torch.npu.Stream + + +TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) + +# use int value instead of ReduceOp.SUM to support torch compile +REDUCE_OP_SUM = int(torch.distributed.ReduceOp.SUM) + + +def _split_tensor_dict( + tensor_dict: Dict[str, Union[torch.Tensor, Any]] +) -> Tuple[List[Tuple[str, Any]], List[torch.Tensor]]: + """Split the tensor dictionary into two parts: + 1. A list of (key, value) pairs. If the value is a tensor, it is replaced + by its metadata. + 2. A list of tensors. + """ + metadata_list: List[Tuple[str, Any]] = [] + tensor_list: List[torch.Tensor] = [] + for key, value in tensor_dict.items(): + if isinstance(value, torch.Tensor): + # Note: we cannot use `value.device` here, + # because it contains not only the device type but also the device + # index (e.g. "cuda:0"). We only need the device type. + # receiving side will set the device index. + device = value.device.type + metadata_list.append( + (key, TensorMetadata(device, value.dtype, value.size())) + ) + tensor_list.append(value) + else: + metadata_list.append((key, value)) + return metadata_list, tensor_list + + +_group_name_counter: Dict[str, int] = {} + + +def _get_unique_name(name: str) -> str: + """Get a unique name for the group. + Example: + _get_unique_name("tp") -> "tp:0" + _get_unique_name("tp") -> "tp:1" + """ + if name not in _group_name_counter: + _group_name_counter[name] = 0 + newname = f"{name}:{_group_name_counter[name]}" + _group_name_counter[name] += 1 + return newname + + +_groups: Dict[str, Callable[[], Optional["GroupCoordinator"]]] = {} + + +def _register_group(group: "GroupCoordinator") -> None: + _groups[group.unique_name] = weakref.ref(group) + + +if supports_custom_op(): + + def inplace_all_reduce(tensor: torch.Tensor, group_name: str) -> None: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + group._all_reduce_in_place(tensor) + + def inplace_all_reduce_fake(tensor: torch.Tensor, group_name: str) -> None: + return + + direct_register_custom_op( + op_name="inplace_all_reduce", + op_func=inplace_all_reduce, + mutates_args=["tensor"], + fake_impl=inplace_all_reduce_fake, + ) + + def outplace_all_reduce( + tensor: torch.Tensor, group_name: str, outplace_all_reduce_method: str + ) -> torch.Tensor: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + return group._all_reduce_out_place(tensor, outplace_all_reduce_method) + + def outplace_all_reduce_fake( + tensor: torch.Tensor, group_name: str, outplace_all_reduce_method: str + ) -> torch.Tensor: + return torch.empty_like(tensor) + + direct_register_custom_op( + op_name="outplace_all_reduce", + op_func=outplace_all_reduce, + mutates_args=[], + fake_impl=outplace_all_reduce_fake, + ) + + def reg_all_gather_into_tensor( + output: torch.Tensor, input: torch.Tensor, group_name: str + ) -> None: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + group._all_gather_into_tensor(output, input) + + def reg_all_gather_into_tensor_fake( + output: torch.Tensor, input: torch.Tensor, group_name: str + ) -> None: + pass + + direct_register_custom_op( + op_name="reg_all_gather_into_tensor", + op_func=reg_all_gather_into_tensor, + mutates_args=["output"], + fake_impl=reg_all_gather_into_tensor_fake, + ) + + +class GroupCoordinator: + """ + PyTorch ProcessGroup wrapper for a group of processes. + PyTorch ProcessGroup is bound to one specific communication backend, + e.g. NCCL, Gloo, MPI, etc. + GroupCoordinator takes charge of all the communication operations among + the processes in the group. It can route the communication to + a specific implementation (e.g. switch allreduce implementation + based on the tensor size and cuda graph mode). + """ + + # available attributes: + rank: int # global rank + ranks: List[int] # global ranks in the group + world_size: int # size of the group + # difference between `local_rank` and `rank_in_group`: + # if we have a group of size 4 across two nodes: + # Process | Node | Rank | Local Rank | Rank in Group + # 0 | 0 | 0 | 0 | 0 + # 1 | 0 | 1 | 1 | 1 + # 2 | 1 | 2 | 0 | 2 + # 3 | 1 | 3 | 1 | 3 + local_rank: int # local rank used to assign devices + rank_in_group: int # rank inside the group + cpu_group: ProcessGroup # group for CPU communication + device_group: ProcessGroup # group for device communication + use_pynccl: bool # a hint of whether to use PyNccl + use_pymscclpp: bool # a hint of whether to use PyMsccl + use_custom_allreduce: bool # a hint of whether to use CustomAllreduce + use_message_queue_broadcaster: ( + bool # a hint of whether to use message queue broadcaster + ) + # communicators are only created for world size > 1 + pynccl_comm: Optional[Any] # PyNccl communicator + ca_comm: Optional[Any] # Custom allreduce communicator + mq_broadcaster: Optional[Any] # shared memory broadcaster + + def __init__( + self, + group_ranks: List[List[int]], + local_rank: int, + torch_distributed_backend: Union[str, Backend], + use_pynccl: bool, + use_pymscclpp: bool, + use_custom_allreduce: bool, + use_hpu_communicator: bool, + use_xpu_communicator: bool, + use_npu_communicator: bool, + use_message_queue_broadcaster: bool = False, + group_name: Optional[str] = None, + ): + # Set group info + group_name = group_name or "anonymous" + self.unique_name = _get_unique_name(group_name) + _register_group(self) + + # Set rank info + self.rank = torch.distributed.get_rank() + self.local_rank = local_rank + self.device_group = None + self.cpu_group = None + self.local_size = get_int_env_var("LOCAL_SIZE", 0) + + for ranks in group_ranks: + device_group = torch.distributed.new_group( + ranks, backend=torch_distributed_backend + ) + # a group with `gloo` backend, to allow direct coordination between + # processes through the CPU. + cpu_group = torch.distributed.new_group(ranks, backend="gloo") + if self.rank in ranks: + self.ranks = ranks + self.world_size = len(ranks) + self.rank_in_group = ranks.index(self.rank) + self.device_group = device_group + self.cpu_group = cpu_group + + assert self.cpu_group is not None + assert self.device_group is not None + + device_id = 0 if IS_ONE_DEVICE_PER_PROCESS else local_rank + if is_cuda_alike(): + self.device = torch.device(f"cuda:{device_id}") + elif _is_npu: + self.device = torch.device(f"npu:{device_id}") + else: + self.device = torch.device("cpu") + self.device_module = torch.get_device_module(self.device) + + # Import communicators + self.use_pynccl = use_pynccl + self.use_pymscclpp = use_pymscclpp + self.use_custom_allreduce = use_custom_allreduce + self.use_hpu_communicator = use_hpu_communicator + self.use_xpu_communicator = use_xpu_communicator + self.use_npu_communicator = use_npu_communicator + self.use_message_queue_broadcaster = use_message_queue_broadcaster + + # lazy import to avoid documentation build error + from sglang.srt.distributed.device_communicators.custom_all_reduce import ( + CustomAllreduce, + ) + from sglang.srt.distributed.device_communicators.pymscclpp import ( + PyMscclppCommunicator, + ) + from sglang.srt.distributed.device_communicators.pynccl import ( + PyNcclCommunicator, + ) + + if is_hip(): + from sglang.srt.distributed.device_communicators.quick_all_reduce import ( + QuickAllReduce, + qr_rocm_arch_available, + ) + + self.pynccl_comm: Optional[PyNcclCommunicator] = None + if use_pynccl and self.world_size > 1: + self.pynccl_comm = PyNcclCommunicator( + group=self.cpu_group, + device=self.device, + ) + + self.pymscclpp_comm: Optional[PyMscclppCommunicator] = None + if use_pymscclpp and self.world_size > 1: + self.pymscclpp_comm = PyMscclppCommunicator( + group=self.cpu_group, + device=self.device, + ) + + self.ca_comm: Optional[CustomAllreduce] = None + self.qr_comm: Optional[QuickAllReduce] = None + if use_custom_allreduce and self.world_size > 1: + # Initialize a custom fast all-reduce implementation. + try: + self.ca_comm = CustomAllreduce( + group=self.cpu_group, + device=self.device, + ) + except Exception as e: + logger.warning( + f"Setup Custom allreduce failed with {e}. To silence this " + "warning, specify --disable-custom-all-reduce explicitly." + ) + if is_hip(): + try: + # Initialize a custom quick all-reduce implementation for AMD + # when rocm >= gfx942. Quick reduce is designed as a + # complement to custom allreduce. + # Based on quickreduce (https://github.com/mk1-project/quickreduce). + if qr_rocm_arch_available(): + self.qr_comm = QuickAllReduce( + group=self.cpu_group, device=self.device + ) + except Exception as e: + logger.warning(f"Failed to initialize QuickAllReduce: {e}") + + # Create communicator for other hardware backends + from sglang.srt.distributed.device_communicators.hpu_communicator import ( + HpuCommunicator, + ) + from sglang.srt.distributed.device_communicators.npu_communicator import ( + NpuCommunicator, + ) + from sglang.srt.distributed.device_communicators.xpu_communicator import ( + XpuCommunicator, + ) + + self.hpu_communicator: Optional[HpuCommunicator] = None + if use_hpu_communicator and self.world_size > 1: + self.hpu_communicator = HpuCommunicator(group=self.device_group) + + self.xpu_communicator: Optional[XpuCommunicator] = None + if use_xpu_communicator and self.world_size > 1: + self.xpu_communicator = XpuCommunicator(group=self.device_group) + + self.npu_communicator: Optional[NpuCommunicator] = None + if use_npu_communicator and self.world_size > 1: + self.npu_communicator = NpuCommunicator(group=self.device_group) + + # Create message queue + from sglang.srt.distributed.device_communicators.shm_broadcast import ( + MessageQueue, + ) + + self.mq_broadcaster: Optional[MessageQueue] = None + if use_message_queue_broadcaster and self.world_size > 1: + self.mq_broadcaster = MessageQueue.create_from_process_group( + self.cpu_group, 1 << 22, 6 + ) + + def __repr__(self): + return ( + f"ranks={self.ranks} rank={self.rank} local_rank={self.local_rank} use_pynccl={self.use_pynccl} " + f"device_group={self.device_group} cpu_group={self.cpu_group} unique_name={self.unique_name} " + f"world_size={self.world_size} rank_in_group={self.rank_in_group}" + ) + + @property + def first_rank(self): + """Return the global rank of the first process in the group""" + return self.ranks[0] + + @property + def last_rank(self): + """Return the global rank of the last process in the group""" + return self.ranks[-1] + + @property + def is_first_rank(self): + """Return whether the caller is the first process in the group""" + return self.rank == self.first_rank + + @property + def is_last_rank(self): + """Return whether the caller is the last process in the group""" + return self.rank == self.last_rank + + @property + def next_rank(self): + """Return the global rank of the process that follows the caller""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group + 1) % world_size] + + @property + def prev_rank(self): + """Return the global rank of the process that precedes the caller""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group - 1) % world_size] + + @contextmanager + def graph_capture( + self, graph_capture_context: Optional[GraphCaptureContext] = None + ): + if graph_capture_context is None: + stream = self.device_module.Stream() + graph_capture_context = GraphCaptureContext(stream) + else: + stream = graph_capture_context.stream + # We don't need the context of custom quick allreduce because the ipc access + # is already collected in init() and we can capture the quick allreduce directly. + ca_comm = self.ca_comm + maybe_ca_context = nullcontext() if ca_comm is None else ca_comm.capture() + + # ensure all initialization operations complete before attempting to + # capture the graph on another stream + curr_stream = self.device_module.current_stream() + if curr_stream != stream: + stream.wait_stream(curr_stream) + + with self.device_module.stream(stream), maybe_ca_context: + # In graph mode, we have to be very careful about the collective + # operations. The current status is: + # allreduce \ Mode | Eager | Graph | + # -------------------------------------------- + # quick allreduce | enabled | enabled | + # custom allreduce | enabled | enabled | + # PyNccl | disabled| enabled | + # PyMscclpp | disabled| enabled | + # torch.distributed | enabled | disabled| + # + # Note: When custom quick allreduce is enabled, a runtime check + # will be performed. If the tensor size is too small, it will + # automatically fall back to the next available option. + # Note that custom allreduce will have a runtime check, if the + # tensor size is too large, it will fallback to the next + # available option. + # Note that the PyMsccl needs to register the tensor in ahead, + # which will introduce large overhead in the eager case, + # therefore it is only supported in the graph case. + # In summary: We select the appropriate allreduce method for + # each mode based on the algorithm order in the table and + # their usage conditions. + pynccl_comm = self.pynccl_comm + maybe_pynccl_context: Any + if not pynccl_comm: + maybe_pynccl_context = nullcontext() + else: + maybe_pynccl_context = pynccl_comm.change_state( + enable=True, stream=torch.cuda.current_stream() + ) + + pymscclpp_comm = self.pymscclpp_comm + maybe_pymscclpp_context: Any + if not pymscclpp_comm: + maybe_pymscclpp_context = nullcontext() + else: + maybe_pymscclpp_context = pymscclpp_comm.change_state(enable=True) + with maybe_pynccl_context, maybe_pymscclpp_context: + yield graph_capture_context + + def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: + """ + User-facing all-reduce function before we actually call the + all-reduce operation. + + We need this because Dynamo does not support passing an arbitrary + object (`self` in this case) to a custom op. We need to pass the + group name as a string, and then look up the group coordinator from + the group name, dispatch the all-reduce operation to the group + coordinator. + + In addition, PyTorch custom ops do not support mutation or returning + a new tensor in the same op. So we need to figure out if the op is + in-place or out-of-place ahead of time. + """ + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + + if input_.is_cpu: + if is_shm_available(input_.dtype, self.world_size, self.local_size): + torch.ops.sgl_kernel.shm_allreduce(input_, REDUCE_OP_SUM) + else: + torch.distributed.all_reduce(input_, group=self.device_group) + return input_ + + if not supports_custom_op(): + self._all_reduce_in_place(input_) + return input_ + + if self.hpu_communicator is not None and not self.hpu_communicator.disabled: + return self.hpu_communicator.all_reduce(input_) + + if self.xpu_communicator is not None and not self.xpu_communicator.disabled: + return self.xpu_communicator.all_reduce(input_) + + if self.npu_communicator is not None and not self.npu_communicator.disabled: + return self.npu_communicator.all_reduce(input_) + + if ( + self.pynccl_comm is not None + and hasattr(input_, "symmetric_memory") + and input_.symmetric_memory + ): + with self.pynccl_comm.change_state( + enable=True, stream=torch.cuda.current_stream() + ): + self.pynccl_comm.all_reduce(input_) + return input_ + + outplace_all_reduce_method = None + if ( + self.qr_comm is not None + and not self.qr_comm.disabled + and self.qr_comm.should_quick_allreduce(input_) + ): + outplace_all_reduce_method = "qr" + elif ( + self.ca_comm is not None + and not self.ca_comm.disabled + and self.ca_comm.should_custom_ar(input_) + ): + outplace_all_reduce_method = "ca" + elif ( + self.pymscclpp_comm is not None + and not self.pymscclpp_comm.disabled + and self.pymscclpp_comm.should_mscclpp_allreduce(input_) + ): + outplace_all_reduce_method = "pymscclpp" + if outplace_all_reduce_method is not None: + return torch.ops.sglang.outplace_all_reduce( + input_, + group_name=self.unique_name, + outplace_all_reduce_method=outplace_all_reduce_method, + ) + else: + torch.ops.sglang.inplace_all_reduce(input_, group_name=self.unique_name) + return input_ + + def _all_reduce_out_place( + self, input_: torch.Tensor, outplace_all_reduce_method: str + ) -> torch.Tensor: + qr_comm = self.qr_comm + ca_comm = self.ca_comm + pymscclpp_comm = self.pymscclpp_comm + assert any([qr_comm, ca_comm, pymscclpp_comm]) + if outplace_all_reduce_method == "qr": + assert not qr_comm.disabled + out = qr_comm.quick_all_reduce(input_) + elif outplace_all_reduce_method == "ca": + assert not ca_comm.disabled + out = ca_comm.custom_all_reduce(input_) + else: + assert not pymscclpp_comm.disabled + out = pymscclpp_comm.all_reduce(input_) + assert out is not None + return out + + def _all_reduce_in_place(self, input_: torch.Tensor) -> None: + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.all_reduce(input_) + else: + torch.distributed.all_reduce(input_, group=self.device_group) + + def reduce_scatter_tensor( + self, + output: torch.Tensor, + input: torch.Tensor, + ) -> None: + # TODO(ch-wan): support other backends + torch.distributed.reduce_scatter_tensor(output, input, group=self.device_group) + return output + + def reduce_scatter( + self, + output: torch.Tensor, + input_list: List[torch.Tensor], + ) -> None: + # TODO(ch-wan): support other backends + torch.distributed.reduce_scatter(output, input_list, group=self.device_group) + return output + + def reduce_scatterv( + self, + input_: torch.Tensor, + output: Optional[torch.Tensor] = None, + sizes: Optional[List[int]] = None, + ) -> torch.Tensor: + world_size = self.world_size + pynccl_comm = self.pynccl_comm + + with pynccl_comm.change_state(enable=True, stream=torch.cuda.current_stream()): + assert ( + pynccl_comm is not None and not pynccl_comm.disabled + ), "pynccl is required for reduce_scatterv" + + if sizes is not None: + assert len(sizes) == world_size + assert input_.shape[0] == sum(sizes) + chunk_size = sizes[self.rank_in_group] + else: + assert input_.shape[0] % world_size == 0 + chunk_size = input_.shape[0] // world_size + output_shape = (chunk_size,) + input_.shape[1:] + + if output is None: + output = torch.empty( + output_shape, dtype=input_.dtype, device=input_.device + ) + else: + assert output.shape == output_shape + + pynccl_comm.reduce_scatter(output, input_, sizes=sizes) + return output + + def _all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor): + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.all_gather(output, input) + else: + torch.distributed.all_gather_into_tensor( + output, input, group=self.device_group + ) + + def all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor): + if _is_npu or not supports_custom_op(): + self._all_gather_into_tensor(output, input) + else: + torch.ops.sglang.reg_all_gather_into_tensor( + output, input, group_name=self.unique_name + ) + + def all_gather( + self, + input_: torch.Tensor, + dim: int = -1, + output_tensor_list: Optional[List[torch.Tensor]] = None, + ) -> torch.Tensor: + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + if output_tensor_list is not None: + logger.warning( + "Performing in-place all-gather with a group size of 1. " + "This may be unnecessary; consider bypassing it for better efficiency." + ) + output_tensor_list[0].copy_(input_) + return None + else: + return input_ + + if output_tensor_list is not None: + # TODO(ch-wan): support other backends + return torch.distributed.all_gather( + output_tensor_list, input_, group=self.device_group + ) + + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + + # For HPUs, use HPU communicator. + hpu_comm = self.hpu_communicator + if hpu_comm is not None and not hpu_comm.disabled: + return hpu_comm.all_gather(input_, dim) + + # For NPUs, use NPU communicator. + npu_comm = self.npu_communicator + if npu_comm is not None and not npu_comm.disabled: + return npu_comm.all_gather(input_, dim) + + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + input_size = input_.size() + # NOTE: we have to use concat-style all-gather here, + # stack-style all-gather has compatibility issues with + # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 + output_size = (input_size[0] * world_size,) + input_size[1:] + # Allocate output tensor. + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + + # All-gather. + if input_.is_cpu and is_shm_available( + input_.dtype, self.world_size, self.local_size + ): + return torch.ops.sgl_kernel.shm_allgather(input_, dim) + + if input_.is_cpu: + torch.distributed.all_gather_into_tensor( + output_tensor, input_, group=self.device_group + ) + else: + self.all_gather_into_tensor(output_tensor, input_) + + # Reshape + output_tensor = output_tensor.reshape((world_size,) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :] + ) + return output_tensor + + def all_gatherv( + self, + input_: Union[torch.Tensor, List[torch.Tensor]], + sizes: Optional[List[int]] = None, + ) -> Union[torch.Tensor, List[torch.Tensor]]: + """ + Supports varying sizes per rank and input tensor list. + `sizes`: a list of len(world_size) with the number of items per rank to gather. + """ + world_size = self.world_size + pynccl_comm = self.pynccl_comm + + with pynccl_comm.change_state(enable=True, stream=torch.cuda.current_stream()): + assert ( + pynccl_comm is not None and not pynccl_comm.disabled + ), "pynccl is required for all_gatherv" + + def _all_gather_single( + input_: torch.Tensor, sizes: Optional[List[int]] = None + ): + input_size = input_.size() + if sizes is not None: + assert len(sizes) == world_size + assert input_.shape[0] == sizes[self.rank_in_group] + output_size = (sum(sizes),) + input_size[1:] + # 'sizes' is not needed if all inputs in the same group have the same shape + if all(s == sizes[0] for s in sizes): + sizes = None + else: + output_size = (input_size[0] * world_size,) + input_size[1:] + # Allocate output tensor. + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + pynccl_comm.all_gather(output_tensor, input_, sizes=sizes) + return output_tensor + + if isinstance(input_, torch.Tensor): + return _all_gather_single(input_, sizes) + + output_list = [] + pynccl_comm.group_start() + for inp in input_: + output_list.append(_all_gather_single(inp, sizes=sizes)) + pynccl_comm.group_end() + + return output_list + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> Optional[torch.Tensor]: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + if self.xpu_communicator is not None and not self.xpu_communicator.disabled: + return self.xpu_communicator.gather(input_, self.rank_in_group, dst, dim) + # Allocate output tensor. + if self.rank_in_group == dst: + gather_list = [torch.empty_like(input_) for _ in range(world_size)] + else: + gather_list = None + # Gather. + torch.distributed.gather( + input_, gather_list, dst=self.ranks[dst], group=self.device_group + ) + if self.rank_in_group == dst: + output_tensor = torch.cat(gather_list, dim=dim) + else: + output_tensor = None + return output_tensor + + def broadcast(self, input_: torch.Tensor, src: int = 0): + """Broadcast the input tensor. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + # Broadcast. + torch.distributed.broadcast( + input_, src=self.ranks[src], group=self.device_group + ) + return input_ + + def broadcast_object(self, obj: Optional[Any] = None, src: int = 0): + """Broadcast the input object. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj + if self.mq_broadcaster is not None: + assert src == 0, "Message queue broadcaster only supports src=0" + return self.mq_broadcaster.broadcast_object(obj) + if self.rank_in_group == src: + torch.distributed.broadcast_object_list( + [obj], src=self.ranks[src], group=self.cpu_group + ) + return obj + else: + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=self.ranks[src], group=self.cpu_group + ) + return recv[0] + + def broadcast_object_list( + self, obj_list: List[Any], src: int = 0, group: Optional[ProcessGroup] = None + ): + """Broadcast the input object list. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj_list + # Broadcast. + torch.distributed.broadcast_object_list( + obj_list, src=self.ranks[src], group=self.device_group + ) + return obj_list + + def all_gather_object(self, obj: Any) -> List[Any]: + objs = [None] * self.world_size + torch.distributed.all_gather_object(objs, obj, group=self.cpu_group) + return objs + + def send_object(self, obj: Any, dst: int) -> None: + """Send the input object list to the destination rank.""" + """NOTE: `dst` is the local rank of the destination rank.""" + + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + assert dst != self.rank_in_group, ( + "Invalid destination rank. Destination rank is the same " + "as the current rank." + ) + + # Serialize object to tensor and get the size as well + object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8).cuda( + device=torch.cuda.current_device() + ) + + size_tensor = torch.tensor( + [object_tensor.numel()], + dtype=torch.long, + device="cpu", + ) + # Send object size + torch.distributed.send(size_tensor, dst=self.ranks[dst], group=self.cpu_group) + + # Send object + torch.distributed.send( + object_tensor, + dst=self.ranks[dst], + group=self.device_group, + ) + + return None + + def recv_object(self, src: int) -> Any: + """Receive the input object list from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + + assert src < self.world_size, f"Invalid src rank ({src})" + + assert ( + src != self.rank_in_group + ), "Invalid source rank. Source rank is the same as the current rank." + + size_tensor = torch.empty(1, dtype=torch.long, device="cpu") + + # Receive object size + rank_size = torch.distributed.recv( + size_tensor, src=self.ranks[src], group=self.cpu_group + ) + + # Tensor to receive serialized objects into. + object_tensor = torch.empty( # type: ignore[call-overload] + size_tensor.item(), # type: ignore[arg-type] + dtype=torch.uint8, + device=torch.cuda.current_device(), + ) + + rank_object = torch.distributed.recv( + object_tensor, src=self.ranks[src], group=self.device_group + ) + + assert ( + rank_object == rank_size + ), "Received object sender rank does not match the size sender rank." + + obj = pickle.loads(object_tensor.cpu().numpy()) + + return obj + + def broadcast_tensor_dict( + self, + tensor_dict: Optional[Dict[str, Union[torch.Tensor, Any]]] = None, + src: int = 0, + group: Optional[ProcessGroup] = None, + metadata_group: Optional[ProcessGroup] = None, + ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]: + """Broadcast the input tensor dictionary. + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return tensor_dict + + group = self.device_group + metadata_group = self.cpu_group + assert src < self.world_size, f"Invalid src rank ({src})" + + rank_in_group = self.rank_in_group + if rank_in_group == src: + metadata_list: List[Tuple[Any, Any]] = [] + assert isinstance( + tensor_dict, dict + ), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # `metadata_list` lives in CPU memory. + # `broadcast_object_list` has serialization & deserialization, + # all happening on CPU. Therefore, we can use the CPU group. + self.broadcast_object(metadata_list, src=src) + async_handles = [] + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=metadata_group, async_op=True + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=group, async_op=True + ) + async_handles.append(handle) + for async_handle in async_handles: + async_handle.wait() + + else: + metadata_list = self.broadcast_object(None, src=src) + tensor_dict = {} + async_handles = [] + for key, value in metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty( + value.size, dtype=value.dtype, device=value.device + ) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, + src=self.ranks[src], + group=metadata_group, + async_op=True, + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=group, async_op=True + ) + async_handles.append(handle) + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + for async_handle in async_handles: + async_handle.wait() + return tensor_dict + + def send_tensor_dict( + self, + tensor_dict: Dict[str, Union[torch.Tensor, Any]], + dst: Optional[int] = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]: + """Send the input tensor dictionary. + NOTE: `dst` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return tensor_dict + + all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + assert isinstance( + tensor_dict, dict + ), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # Note: While switching to Device-to-Device (D2D) would introduce an extra + # Device-to-Host (D2H) memory copy overhead for serialization, our benchmarks + # show better overall transmission performance with D2D due to: + # 1. Superior D2D transfer bandwidth + # 2. Ability to overlap send and recv operations + # Thus the net performance gain justifies this approach. + self.send_object(metadata_list, dst=dst) + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip sending empty tensors. + continue + + # send-allgather: send only a slice, then do allgather. + if all_gather_group is not None and tensor.numel() % all_gather_size == 0: + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.send( + tensor, dst=self.ranks[dst], group=metadata_group + ) + else: + # use group for GPU tensors + torch.distributed.send(tensor, dst=self.ranks[dst], group=group) + return None + + def recv_tensor_dict( + self, + src: Optional[int] = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]: + """Recv the input tensor dictionary. + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return None + + all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if src is None: + src = (self.rank_in_group - 1) % self.world_size + assert src < self.world_size, f"Invalid src rank ({src})" + + recv_metadata_list = self.recv_object(src=src) + tensor_dict: Dict[str, Any] = {} + for key, value in recv_metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty(value.size, dtype=value.dtype, device=value.device) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + + # send-allgather: send only a slice, then do allgather. + use_all_gather = ( + all_gather_group is not None + and tensor.numel() % all_gather_size == 0 + ) + + if use_all_gather: + orig_shape = tensor.shape + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.recv( + tensor, src=self.ranks[src], group=metadata_group + ) + else: + # use group for GPU tensors + torch.distributed.recv(tensor, src=self.ranks[src], group=group) + if use_all_gather: + # do the allgather + tensor = all_gather_group.all_gather(tensor, dim=0) # type: ignore + tensor = tensor.reshape(orig_shape) + + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + return tensor_dict + + def barrier(self): + """Barrier synchronization among the group. + NOTE: don't use `device_group` here! `barrier` in NCCL is + terrible because it is internally a broadcast operation with + secretly created GPU tensors. It is easy to mess up the current + device. Use the CPU group instead. + """ + torch.distributed.barrier(group=self.cpu_group) + + def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way""" + """NOTE: `dst` is the local rank of the destination rank.""" + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.send(tensor, dst) + else: + torch.distributed.send(tensor, self.ranks[dst], self.device_group) + + def recv( + self, size: torch.Size, dtype: torch.dtype, src: Optional[int] = None + ) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + if src is None: + src = (self.rank_in_group - 1) % self.world_size + + tensor = torch.empty(size, dtype=dtype, device=self.device) + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.recv(tensor, src) + else: + torch.distributed.recv(tensor, self.ranks[src], self.device_group) + return tensor + + def destroy(self): + if self.device_group is not None: + torch.distributed.destroy_process_group(self.device_group) + self.device_group = None + if self.cpu_group is not None: + torch.distributed.destroy_process_group(self.cpu_group) + self.cpu_group = None + if self.pynccl_comm is not None: + self.pynccl_comm = None + if self.ca_comm is not None: + self.ca_comm = None + if self.mq_broadcaster is not None: + self.mq_broadcaster = None + + +_WORLD: Optional[GroupCoordinator] = None + + +def get_world_group() -> GroupCoordinator: + assert _WORLD is not None, "world group is not initialized" + return _WORLD + + +def init_world_group( + ranks: List[int], local_rank: int, backend: str +) -> GroupCoordinator: + return GroupCoordinator( + group_ranks=[ranks], + local_rank=local_rank, + torch_distributed_backend=backend, + use_pynccl=False, + use_pymscclpp=False, + use_custom_allreduce=False, + use_hpu_communicator=False, + use_xpu_communicator=False, + use_npu_communicator=False, + group_name="world", + ) + + +def init_model_parallel_group( + group_ranks: List[List[int]], + local_rank: int, + backend: str, + use_custom_allreduce: Optional[bool] = None, + use_message_queue_broadcaster: bool = False, + group_name: Optional[str] = None, + use_mscclpp_allreduce: Optional[bool] = None, +) -> GroupCoordinator: + if use_custom_allreduce is None: + use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE + if use_mscclpp_allreduce is None: + use_mscclpp_allreduce = _ENABLE_MSCCLPP_ALL_REDUCE + return GroupCoordinator( + group_ranks=group_ranks, + local_rank=local_rank, + torch_distributed_backend=backend, + use_pynccl=not _is_npu, + use_pymscclpp=use_mscclpp_allreduce, + use_custom_allreduce=use_custom_allreduce, + use_hpu_communicator=True, + use_xpu_communicator=True, + use_npu_communicator=True, + use_message_queue_broadcaster=use_message_queue_broadcaster, + group_name=group_name, + ) + + +_TP: Optional[GroupCoordinator] = None + +# duplicate GroupCoordinator for prefill in PD-Multiplexing +_PDMUX_PREFILL_TP_GROUP: Optional[GroupCoordinator] = None + +_ENABLE_PDMUX_P_TP: bool = False + + +def set_pdmux_status(enable_prefill_multiplexing: bool): + global _ENABLE_PDMUX_P_TP + _ENABLE_PDMUX_P_TP = enable_prefill_multiplexing + + +def get_tp_group() -> GroupCoordinator: + if _ENABLE_PDMUX_P_TP: + assert ( + _PDMUX_PREFILL_TP_GROUP is not None + ), "tensor model parallel group for PD-Multiplexing Prefill is not initialized" + return _PDMUX_PREFILL_TP_GROUP + assert _TP is not None, "tensor model parallel group is not initialized" + return _TP + + +_MOE_EP: Optional[GroupCoordinator] = None +_MOE_TP: Optional[GroupCoordinator] = None + + +def get_moe_ep_group() -> GroupCoordinator: + assert _MOE_EP is not None, "expert model parallel group is not initialized" + return _MOE_EP + + +def get_moe_tp_group() -> GroupCoordinator: + assert _MOE_TP is not None, "expert model parallel group is not initialized" + return _MOE_TP + + +# kept for backward compatibility +get_tensor_model_parallel_group = get_tp_group + +_PP: Optional[GroupCoordinator] = None + + +def get_pp_group() -> GroupCoordinator: + assert _PP is not None, "pipeline model parallel group is not initialized" + return _PP + + +# kept for backward compatibility +get_pipeline_model_parallel_group = get_pp_group + + +@contextmanager +def graph_capture(): + """ + `graph_capture` is a context manager which should surround the code that + is capturing the CUDA graph. Its main purpose is to ensure that the + some operations will be run after the graph is captured, before the graph + is replayed. It returns a `GraphCaptureContext` object which contains the + necessary data for the graph capture. Currently, it only contains the + stream that the graph capture is running on. This stream is set to the + current CUDA stream when the context manager is entered and reset to the + default stream when the context manager is exited. This is to ensure that + the graph capture is running on a separate stream from the default stream, + in order to explicitly distinguish the kernels to capture + from other kernels possibly launched on background in the default stream. + """ + with get_tp_group().graph_capture() as context, get_pp_group().graph_capture( + context + ): + yield context + + +logger = logging.getLogger(__name__) + +_ENABLE_CUSTOM_ALL_REDUCE = True +_ENABLE_MSCCLPP_ALL_REDUCE = False + + +def set_custom_all_reduce(enable: bool): + global _ENABLE_CUSTOM_ALL_REDUCE + _ENABLE_CUSTOM_ALL_REDUCE = enable + + +def set_mscclpp_all_reduce(enable: bool): + global _ENABLE_MSCCLPP_ALL_REDUCE + _ENABLE_MSCCLPP_ALL_REDUCE = enable + + +def init_distributed_environment( + world_size: int = -1, + rank: int = -1, + distributed_init_method: str = "env://", + local_rank: int = -1, + backend: str = "nccl", + timeout: Optional[int] = None, +): + logger.debug( + "world_size=%d rank=%d local_rank=%d " "distributed_init_method=%s backend=%s", + world_size, + rank, + local_rank, + distributed_init_method, + backend, + ) + if not torch.distributed.is_initialized(): + assert distributed_init_method is not None, ( + "distributed_init_method must be provided when initializing " + "distributed environment" + ) + if timeout is not None: + assert isinstance(timeout, (int)), "timeout must be a number" + assert timeout > 0, "timeout must be positive" + timeout = timedelta(seconds=timeout) + + # this backend is used for WORLD + torch.distributed.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=world_size, + rank=rank, + timeout=timeout, + ) + + # set the local rank + # local_rank is not available in torch ProcessGroup, + # see https://github.com/pytorch/pytorch/issues/122816 + if local_rank == -1: + # local rank not set, this usually happens in single-node + # setting, where we can use rank as local rank + if distributed_init_method == "env://": + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + else: + local_rank = rank + global _WORLD + if _WORLD is None: + ranks = list(range(torch.distributed.get_world_size())) + _WORLD = init_world_group(ranks, local_rank, backend) + else: + assert ( + _WORLD.world_size == torch.distributed.get_world_size() + ), "world group already initialized with a different world size" + + +def initialize_model_parallel( + tensor_model_parallel_size: int = 1, + expert_model_parallel_size: int = 1, + pipeline_model_parallel_size: int = 1, + backend: Optional[str] = None, + duplicate_tp_group: bool = False, +) -> None: + """ + Initialize model parallel groups. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model + parallelism. + pipeline_model_parallel_size: number of GPUs used for pipeline model + parallelism. + + Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we + use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize + the model pipeline. The present function will + create 4 tensor model-parallel groups and 2 pipeline model-parallel groups: + 4 tensor model-parallel groups: + [g0, g1], [g2, g3], [g4, g5], [g6, g7] + 2 pipeline model-parallel groups: + [g0, g2, g4, g6], [g1, g3, g5, g7] + Note that for efficiency, the caller should make sure adjacent ranks + are on the same DGX box. For example if we are using 2 DGX-1 boxes + with a total of 16 GPUs, rank 0 to 7 belong to the first box and + ranks 8 to 15 belong to the second box. + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend(get_world_group().device_group) + + if world_size != tensor_model_parallel_size * pipeline_model_parallel_size: + raise RuntimeError( + f"world_size ({world_size}) is not equal to " + f"tensor_model_parallel_size ({tensor_model_parallel_size}) x " + f"pipeline_model_parallel_size ({pipeline_model_parallel_size})" + ) + + # Build the tensor model-parallel groups. + num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size + global _TP + assert _TP is None, "tensor model parallel group is already initialized" + group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) + ) + group_ranks.append(ranks) + + # message queue broadcaster is only used in tensor model parallel group + _TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=get_bool_env_var( + "SGLANG_USE_MESSAGE_QUEUE_BROADCASTER", "true" + ), + group_name="tp", + ) + + if duplicate_tp_group: + global _PDMUX_PREFILL_TP_GROUP + assert ( + _PDMUX_PREFILL_TP_GROUP is None + ), "tensor model parallel group for PD-Multiplexing Prefill is already initialized" + _PDMUX_PREFILL_TP_GROUP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=get_bool_env_var( + "SGLANG_USE_MESSAGE_QUEUE_BROADCASTER", "true" + ), + group_name="pdmux_prefill_tp", + ) + _TP.pynccl_comm.disabled = False + _PDMUX_PREFILL_TP_GROUP.pynccl_comm.disabled = False + + moe_ep_size = expert_model_parallel_size + moe_tp_size = tensor_model_parallel_size // moe_ep_size + + global _MOE_EP + assert _MOE_EP is None, "expert model parallel group is already initialized" + + if moe_ep_size == tensor_model_parallel_size: + _MOE_EP = _TP + else: + # TODO(ch-wan): use split_group to save memory + group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + for j in range(moe_tp_size): + st = i * tensor_model_parallel_size + j + en = (i + 1) * tensor_model_parallel_size + j + ranks = list(range(st, en, moe_tp_size)) + group_ranks.append(ranks) + _MOE_EP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="moe_ep", + ) + + global _MOE_TP + assert _MOE_TP is None, "expert model parallel group is already initialized" + + if moe_tp_size == tensor_model_parallel_size: + _MOE_TP = _TP + else: + # TODO(ch-wan): use split_group to save memory + group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + for j in range(moe_ep_size): + st = i * tensor_model_parallel_size + j * moe_tp_size + en = i * tensor_model_parallel_size + (j + 1) * moe_tp_size + ranks = list(range(st, en)) + group_ranks.append(ranks) + _MOE_TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="moe_tp", + ) + + # Build the pipeline model-parallel groups. + num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size + global _PP + assert _PP is None, "pipeline model parallel group is already initialized" + group_ranks = [] + for i in range(num_pipeline_model_parallel_groups): + ranks = list(range(i, world_size, num_pipeline_model_parallel_groups)) + group_ranks.append(ranks) + # pipeline parallel does not need custom allreduce + _PP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_custom_allreduce=False, + group_name="pp", + ) + + +def ensure_model_parallel_initialized( + tensor_model_parallel_size: int, + expert_model_parallel_size: int, + pipeline_model_parallel_size: int, + backend: Optional[str] = None, +) -> None: + """Helper to initialize model parallel groups if they are not initialized, + or ensure tensor-parallel and pipeline-parallel sizes are equal to expected + values if the model parallel groups are initialized. + """ + backend = backend or torch.distributed.get_backend(get_world_group().device_group) + if not model_parallel_is_initialized(): + initialize_model_parallel( + tensor_model_parallel_size, + expert_model_parallel_size, + pipeline_model_parallel_size, + backend, + ) + return + + assert get_tensor_model_parallel_world_size() == tensor_model_parallel_size, ( + "tensor parallel group already initialized, but of unexpected size: " + f"{get_tensor_model_parallel_world_size()=} vs. " + f"{tensor_model_parallel_size=}" + ) + pp_world_size = get_pp_group().world_size + assert pp_world_size == pipeline_model_parallel_size, ( + "pipeline parallel group already initialized, but of unexpected size: " + f"{pp_world_size=} vs. " + f"{pipeline_model_parallel_size=}" + ) + + +def model_parallel_is_initialized(): + """Check if tensor and pipeline parallel groups are initialized.""" + return _TP is not None and _PP is not None + + +_TP_STATE_PATCHED = False + + +@contextmanager +def patch_tensor_parallel_group(tp_group: GroupCoordinator): + """Patch the tp group temporarily until this function ends. + + This method is for draft workers of speculative decoding to run draft model + with different tp degree from that of target model workers. + + Args: + tp_group (GroupCoordinator): the tp group coordinator + """ + global _TP_STATE_PATCHED + assert not _TP_STATE_PATCHED, "Should not call when it's already patched" + + _TP_STATE_PATCHED = True + old_tp_group = get_tp_group() + global _TP + _TP = tp_group + try: + yield + finally: + # restore the original state + _TP_STATE_PATCHED = False + _TP = old_tp_group + + +def get_world_size(): + """Return world size for the world group.""" + return get_world_group().world_size + + +def get_world_rank(): + """Return my rank for the world group.""" + return get_world_group().rank_in_group + + +def get_tensor_model_parallel_world_size(): + """Return world size for the tensor model parallel group.""" + return get_tp_group().world_size + + +def get_tensor_model_parallel_rank(): + """Return my rank for the tensor model parallel group.""" + return get_tp_group().rank_in_group + + +def get_pipeline_model_parallel_world_size(): + """Return world size for the pipeline model parallel group.""" + return get_pp_group().world_size + + +def get_pipeline_model_parallel_rank(): + """Return my rank for the pipeline model parallel group.""" + return get_pp_group().rank_in_group + + +def get_moe_expert_parallel_world_size(): + """Return world size for the moe expert parallel group.""" + return get_moe_ep_group().world_size + + +def get_moe_expert_parallel_rank(): + """Return my rank for the moe expert parallel group.""" + return get_moe_ep_group().rank_in_group + + +def get_moe_tensor_parallel_world_size(): + """Return world size for the moe tensor parallel group.""" + return get_moe_tp_group().world_size + + +def get_moe_tensor_parallel_rank(): + """Return my rank for the moe tensor parallel group.""" + return get_moe_tp_group().rank_in_group + + +def destroy_model_parallel(): + """Set the groups to none and destroy them.""" + global _TP + if _TP: + _TP.destroy() + _TP = None + + global _PP + if _PP: + _PP.destroy() + _PP = None + + +def destroy_distributed_environment(): + global _WORLD + if _WORLD: + _WORLD.destroy() + _WORLD = None + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def cleanup_dist_env_and_memory(shutdown_ray: bool = False): + destroy_model_parallel() + destroy_distributed_environment() + with contextlib.suppress(AssertionError): + torch.distributed.destroy_process_group() + if shutdown_ray: + import ray # Lazy import Ray + + ray.shutdown() + gc.collect() + if not _is_cpu: + if hasattr(torch, "cuda") and torch.cuda.is_available(): + torch.cuda.empty_cache() + if hasattr(torch._C, "_host_emptyCache"): + torch._C._host_emptyCache() + else: + logger.warning( + "torch._C._host_emptyCache() only available in Pytorch >=2.5" + ) + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() + elif hasattr(torch, "npu") and torch.npu.is_available(): + torch.npu.empty_cache() + + +def in_the_same_node_as(pg: ProcessGroup, source_rank: int = 0) -> List[bool]: + """ + This is a collective operation that returns if each rank is in the same node + as the source rank. It tests if processes are attached to the same + memory system (shared access to shared memory). + """ + assert ( + torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL + ), "in_the_same_node_as should be tested with a non-NCCL group." + # local rank inside the group + rank = torch.distributed.get_rank(group=pg) + world_size = torch.distributed.get_world_size(group=pg) + + # local tensor in each process to store the result + is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32) + + # global ranks of the processes in the group + ranks = torch.distributed.get_process_group_ranks(pg) + + magic_message = b"magic_message" + shm = None + + try: + with contextlib.suppress(OSError): + if rank == source_rank: + # create a shared memory segment + shm = shared_memory.SharedMemory(create=True, size=128) + shm.buf[: len(magic_message)] = magic_message + torch.distributed.broadcast_object_list( + [shm.name], src=ranks[source_rank], group=pg + ) + is_in_the_same_node[rank] = 1 + else: + # try to open the shared memory segment + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=ranks[source_rank], group=pg + ) + name = recv[0] + # fix to https://stackoverflow.com/q/62748654/9191338 + # Python incorrectly tracks shared memory even if it is not + # created by the process. The following patch is a workaround. + with patch( + "multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None, + ): + shm = shared_memory.SharedMemory(name=name) + if shm.buf[: len(magic_message)] == magic_message: + is_in_the_same_node[rank] = 1 + except Exception as e: + logger.error("Error ignored in is_in_the_same_node: %s", e) + finally: + if shm: + shm.close() + + torch.distributed.barrier(group=pg) + + # clean up the shared memory segment + with contextlib.suppress(OSError): + if rank == source_rank and shm: + shm.unlink() + torch.distributed.all_reduce(is_in_the_same_node, group=pg) + + return [x == 1 for x in is_in_the_same_node.tolist()] + + +vllm_get_pp_group = None +vllm_get_tp_group = None +vllm_get_world_group = None + + +def monkey_patch_vllm_parallel_state(reverse: bool = False): + try: + import vllm.distributed.parallel_state as vllm_parrlel_state + except ImportError: + return + + global vllm_get_pp_group, vllm_get_tp_group, vllm_get_world_group + if vllm_get_pp_group is None: + vllm_get_pp_group = vllm_parrlel_state.get_pp_group + vllm_get_tp_group = vllm_parrlel_state.get_tp_group + vllm_get_world_group = vllm_parrlel_state.get_world_group + if reverse: + setattr(vllm_parrlel_state, "get_pp_group", vllm_get_pp_group) + setattr(vllm_parrlel_state, "get_tp_group", vllm_get_tp_group) + setattr(vllm_parrlel_state, "get_world_group", vllm_get_world_group) + else: + setattr(vllm_parrlel_state, "get_pp_group", get_pp_group) + setattr(vllm_parrlel_state, "get_tp_group", get_tp_group) + setattr(vllm_parrlel_state, "get_world_group", get_world_group) diff --git a/python/sglang/srt/distributed/utils.py b/python/sglang/srt/distributed/utils.py new file mode 100644 index 000000000..bfe54b9d4 --- /dev/null +++ b/python/sglang/srt/distributed/utils.py @@ -0,0 +1,224 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/utils.py + +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +import dataclasses +import logging +import os +import pickle +import time +from collections import deque +from typing import Any, Deque, Dict, Optional, Sequence, Tuple + +import torch +from torch.distributed import TCPStore + +logger = logging.getLogger(__name__) + + +def ensure_divisibility(numerator, denominator): + """Ensure that numerator is divisible by the denominator.""" + assert numerator % denominator == 0, "{} is not divisible by {}".format( + numerator, denominator + ) + + +def divide(numerator, denominator): + """Ensure that numerator is divisible by the denominator and return + the division value.""" + ensure_divisibility(numerator, denominator) + return numerator // denominator + + +def split_tensor_along_last_dim( + tensor: torch.Tensor, + num_partitions: int, + contiguous_split_chunks: bool = False, +) -> Sequence[torch.Tensor]: + """Split a tensor along its last dimension. + + Arguments: + tensor: input tensor. + num_partitions: number of partitions to split the tensor + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + + Returns: + A list of Tensors + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + last_dim_size = divide(tensor.size()[last_dim], num_partitions) + # Split. + tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) + # NOTE: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tensor_list + + +def get_pp_indices( + num_hidden_layers: int, pp_rank: int, pp_size: int +) -> Tuple[int, int]: + """Try to evenly distribute layers across partitions. + If the number of layers is not divisible by the number of partitions, + the last partition will have the remaining layers. + """ + # partition_list_str can be set to None in sglang + partition_list_str = os.getenv("SGLANG_PP_LAYER_PARTITION", None) + if partition_list_str is not None: + try: + partitions = [int(layer) for layer in partition_list_str.split(",")] + except ValueError as err: + raise ValueError( + "Invalid partition string: {}".format(partition_list_str) + ) from err + if len(partitions) != pp_size: + raise ValueError(f"{len(partitions)=} does not match {pp_size=}.") + if sum(partitions) != num_hidden_layers: + raise ValueError(f"{sum(partitions)=} does not match {num_hidden_layers=}.") + start_layer = sum(partitions[:pp_rank]) + end_layer = start_layer + partitions[pp_rank] + else: + layers_per_partition = num_hidden_layers // pp_size + start_layer = pp_rank * layers_per_partition + end_layer = start_layer + layers_per_partition + + if pp_rank == pp_size - 1: + end_layer = num_hidden_layers + + return (start_layer, end_layer) + + +@dataclasses.dataclass +class StatelessProcessGroup: + """A dataclass to hold a metadata store, and the rank, world_size of the + group. Only use it to communicate metadata between processes. + For data-plane communication, create NCCL-related objects. + """ + + rank: int + world_size: int + store: torch._C._distributed_c10d.Store + data_expiration_seconds: int = 3600 # 1 hour + + # dst rank -> counter + send_dst_counter: Dict[int, int] = dataclasses.field(default_factory=dict) + # src rank -> counter + recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict) + broadcast_send_counter: int = 0 + broadcast_recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict) + + # A deque to store the data entries, with key and timestamp. + entries: Deque[Tuple[str, float]] = dataclasses.field(default_factory=deque) + + def __post_init__(self): + assert self.rank < self.world_size + self.send_dst_counter = {i: 0 for i in range(self.world_size)} + self.recv_src_counter = {i: 0 for i in range(self.world_size)} + self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)} + + def send_obj(self, obj: Any, dst: int): + """Send an object to a destination rank.""" + self.expire_data() + key = f"send_to/{dst}/{self.send_dst_counter[dst]}" + self.store.set(key, pickle.dumps(obj)) + self.send_dst_counter[dst] += 1 + self.entries.append((key, time.perf_counter())) + + def expire_data(self): + """Expire data that is older than `data_expiration_seconds` seconds.""" + while self.entries: + # check the oldest entry + key, timestamp = self.entries[0] + if time.perf_counter() - timestamp > self.data_expiration_seconds: + self.store.delete_key(key) + self.entries.popleft() + else: + break + + def recv_obj(self, src: int) -> Any: + """Receive an object from a source rank.""" + obj = pickle.loads( + self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}") + ) + self.recv_src_counter[src] += 1 + return obj + + def broadcast_obj(self, obj: Optional[Any], src: int) -> Any: + """Broadcast an object from a source rank to all other ranks. + It does not clean up after all ranks have received the object. + Use it for limited times, e.g., for initialization. + """ + if self.rank == src: + self.expire_data() + key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}" + self.store.set(key, pickle.dumps(obj)) + self.broadcast_send_counter += 1 + self.entries.append((key, time.perf_counter())) + return obj + else: + key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}" + recv_obj = pickle.loads(self.store.get(key)) + self.broadcast_recv_src_counter[src] += 1 + return recv_obj + + def all_gather_obj(self, obj: Any) -> list[Any]: + """All gather an object from all ranks.""" + gathered_objs = [] + for i in range(self.world_size): + if i == self.rank: + gathered_objs.append(obj) + self.broadcast_obj(obj, src=self.rank) + else: + recv_obj = self.broadcast_obj(None, src=i) + gathered_objs.append(recv_obj) + return gathered_objs + + def barrier(self): + """A barrier to synchronize all ranks.""" + for i in range(self.world_size): + if i == self.rank: + self.broadcast_obj(None, src=self.rank) + else: + self.broadcast_obj(None, src=i) + + @staticmethod + def create( + host: str, + port: int, + rank: int, + world_size: int, + data_expiration_seconds: int = 3600, + ) -> "StatelessProcessGroup": + """A replacement for `torch.distributed.init_process_group` that does not + pollute the global state. + + If we have process A and process B called `torch.distributed.init_process_group` + to form a group, and then we want to form another group with process A, B, C, + D, it is not possible in PyTorch, because process A and process B have already + formed a group, and process C and process D cannot join that group. This + function is a workaround for this issue. + + `torch.distributed.init_process_group` is a global call, while this function + is a stateless call. It will return a `StatelessProcessGroup` object that can be + used for exchanging metadata. With this function, process A and process B + can call `StatelessProcessGroup.create` to form a group, and then process A, B, + C, and D can call `StatelessProcessGroup.create` to form another group. + """ # noqa + store = TCPStore( + host_name=host, + port=port, + world_size=world_size, + is_master=(rank == 0), + ) + + return StatelessProcessGroup( + rank=rank, + world_size=world_size, + store=store, + data_expiration_seconds=data_expiration_seconds, + ) diff --git a/python/sglang/srt/entrypoints/EngineBase.py b/python/sglang/srt/entrypoints/EngineBase.py new file mode 100644 index 000000000..42ecb12aa --- /dev/null +++ b/python/sglang/srt/entrypoints/EngineBase.py @@ -0,0 +1,72 @@ +from abc import ABC, abstractmethod +from typing import Dict, Iterator, List, Optional, Tuple, Union + +import torch + + +class EngineBase(ABC): + """ + Abstract base class for engine interfaces that support generation, weight updating, and memory control. + This base class provides a unified API for both HTTP-based engines and engines. + """ + + @abstractmethod + def generate( + self, + prompt: Optional[Union[List[str], str]] = None, + sampling_params: Optional[Union[List[Dict], Dict]] = None, + input_ids: Optional[Union[List[List[int]], List[int]]] = None, + image_data: Optional[Union[List[str], str]] = None, + return_logprob: Optional[Union[List[bool], bool]] = False, + logprob_start_len: Optional[Union[List[int], int]] = None, + top_logprobs_num: Optional[Union[List[int], int]] = None, + token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None, + lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None, + custom_logit_processor: Optional[Union[List[str], str]] = None, + return_hidden_states: Optional[bool] = None, + stream: Optional[bool] = None, + bootstrap_host: Optional[Union[List[str], str]] = None, + bootstrap_port: Optional[Union[List[int], int]] = None, + bootstrap_room: Optional[Union[List[int], int]] = None, + data_parallel_rank: Optional[int] = None, + ) -> Union[Dict, Iterator[Dict]]: + """Generate outputs based on given inputs.""" + pass + + @abstractmethod + def flush_cache(self): + """Flush the cache of the engine.""" + pass + + @abstractmethod + def update_weights_from_tensor( + self, + named_tensors: List[Tuple[str, torch.Tensor]], + load_format: Optional[str] = None, + flush_cache: bool = True, + ): + """Update model weights with in-memory tensor data.""" + pass + + def load_lora_adapter(self, lora_name: str, lora_path: str): + """Load a new LoRA adapter without re-launching the engine.""" + pass + + def unload_lora_adapter(self, lora_name: str): + """Unload a LoRA adapter without re-launching the engine.""" + pass + + @abstractmethod + def release_memory_occupation(self): + """Release GPU memory occupation temporarily.""" + pass + + @abstractmethod + def resume_memory_occupation(self): + """Resume GPU memory occupation which is previously released.""" + pass + + @abstractmethod + def shutdown(self): + """Shutdown the engine and clean up resources.""" + pass diff --git a/python/sglang/srt/entrypoints/context.py b/python/sglang/srt/entrypoints/context.py new file mode 100644 index 000000000..66f58200f --- /dev/null +++ b/python/sglang/srt/entrypoints/context.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copied from vLLM +import json +import logging +from abc import ABC, abstractmethod +from typing import Union + +logger = logging.getLogger(__name__) + +try: + from mcp import ClientSession +except ImportError as e: + mcp = e + +from openai_harmony import Author, Message, Role, StreamState, TextContent + +from sglang.srt.entrypoints.harmony_utils import ( + get_encoding, + get_streamable_parser_for_assistant, + render_for_completion, +) +from sglang.srt.entrypoints.tool import Tool + + +class ConversationContext(ABC): + + @abstractmethod + def append_output(self, output) -> None: + pass + + @abstractmethod + async def call_tool(self) -> list[Message]: + pass + + @abstractmethod + def need_builtin_tool_call(self) -> bool: + pass + + @abstractmethod + def render_for_completion(self) -> list[int]: + pass + + +class SimpleContext(ConversationContext): + + def __init__(self): + self.last_output = None + + def append_output(self, output) -> None: + self.last_output = output + + def need_builtin_tool_call(self) -> bool: + return False + + async def call_tool(self) -> list[Message]: + raise NotImplementedError("Should not be called.") + + def render_for_completion(self) -> list[int]: + raise NotImplementedError("Should not be called.") + + +class HarmonyContext(ConversationContext): + + def __init__( + self, + messages: list, + tool_sessions: dict[str, Union["ClientSession", Tool]], + ): + # TODO: Remove the hack of Union[ClientSession, Tool] by using MCP + # when demo. + self._messages = messages + self.tool_sessions = tool_sessions + + self.parser = get_streamable_parser_for_assistant() + self.num_init_messages = len(messages) + # TODO + self.num_prompt_tokens = 0 + self.num_cached_tokens = 0 + self.num_output_tokens = 0 + self.num_reasoning_tokens = 0 + + def append_output(self, output) -> None: + if isinstance(output, dict) and "output_ids" in output: + output_token_ids = output["output_ids"] + + # TODO: REMOVE here: + # Very hacky, find the first occurrence of token 200006 and cut from there + try: + start_index = output_token_ids.index(200006) + output_token_ids = output_token_ids[start_index:] + except ValueError: + pass + + for token_id in output_token_ids: + self.parser.process(token_id) + output_msgs = self.parser.messages + + meta_info = output["meta_info"] + + if isinstance(meta_info, dict): + if "prompt_token_ids" in meta_info: + self.num_prompt_tokens = meta_info["prompt_tokens"] + if "cached_tokens" in meta_info: + self.num_cached_tokens = meta_info["cached_tokens"] + if "completion_tokens" in meta_info: + self.num_output_tokens += meta_info["completion_tokens"] + + else: + output_msgs = output + + self._messages.extend(output_msgs) + + @property + def messages(self) -> list: + return self._messages + + def need_builtin_tool_call(self) -> bool: + if not self.messages: + return False + last_msg = self.messages[-1] + recipient = last_msg.recipient + return recipient is not None and ( + recipient.startswith("browser.") or recipient.startswith("python") + ) + + async def call_tool(self) -> list[Message]: + if not self.messages: + return [] + last_msg = self.messages[-1] + recipient = last_msg.recipient + if recipient is not None: + if recipient.startswith("browser."): + return await self.call_search_tool( + self.tool_sessions["browser"], last_msg + ) + elif recipient.startswith("python"): + return await self.call_python_tool( + self.tool_sessions["python"], last_msg + ) + raise ValueError("No tool call found") + + def render_for_completion(self) -> list[int]: + return render_for_completion(self.messages) + + async def call_search_tool( + self, tool_session: Union["ClientSession", Tool], last_msg: Message + ) -> list[Message]: + if isinstance(tool_session, Tool): + return await tool_session.get_result(self) + tool_name = last_msg.recipient.split(".")[1] + args = json.loads(last_msg.content[0].text) + result = await tool_session.call_tool(tool_name, args) + result_str = result.content[0].text + content = TextContent(text=result_str) + author = Author(role=Role.TOOL, name=last_msg.recipient) + return [Message(author=author, content=[content], recipient=Role.ASSISTANT)] + + async def call_python_tool( + self, tool_session: Union["ClientSession", Tool], last_msg: Message + ) -> list[Message]: + if isinstance(tool_session, Tool): + return await tool_session.get_result(self) + param = { + "code": last_msg.content[0].text, + } + result = await tool_session.call_tool("python", param) + result_str = result.content[0].text + + content = TextContent(text=result_str) + author = Author(role=Role.TOOL, name="python") + + return [ + Message( + author=author, + content=[content], + channel=last_msg.channel, + recipient=Role.ASSISTANT, + ) + ] + + +class StreamingHarmonyContext(HarmonyContext): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.last_output = None + + self.parser = get_streamable_parser_for_assistant() + self.encoding = get_encoding() + self.last_tok = None + + @property + def messages(self) -> list: + return self.parser.messages + + def append_output(self, output) -> None: + if isinstance(output, dict) and "output_ids" in output: + # RequestOutput from SGLang with outputs + output_token_ids = output["output_ids"] + + # TODO: REMOVE here: + # Very hacky, find the first occurrence of token 200006 and cut from there + # Find the first occurrence of token 200006 and cut from there + try: + start_index = output_token_ids.index(200006) + output_token_ids = output_token_ids[start_index:] + except ValueError: + pass + + for token_id in output_token_ids: + self.parser.process(token_id) + + else: + # Handle the case of tool output in direct message format + assert len(output) == 1, "Tool output should be a single message" + msg = output[0] + # Sometimes the recipient is not set for tool messages, + # so we set it to "assistant" + if msg.author.role == Role.TOOL and msg.recipient is None: + msg.recipient = "assistant" + toks = self.encoding.render(msg) + for tok in toks: + self.parser.process(tok) + self.last_tok = toks[-1] + + def is_expecting_start(self) -> bool: + return self.parser.state == StreamState.EXPECT_START + + def is_assistant_action_turn(self) -> bool: + return self.last_tok in self.encoding.stop_tokens_for_assistant_actions() + + def render_for_completion(self) -> list[int]: + # now this list of tokens as next turn's starting tokens + # `<|start|>assistant``, + # we need to process them in parser. + rendered_tokens = super().render_for_completion() + + last_n = -1 + to_process = [] + while rendered_tokens[last_n] != self.last_tok: + to_process.append(rendered_tokens[last_n]) + last_n -= 1 + for tok in reversed(to_process): + self.parser.process(tok) + + return rendered_tokens diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py new file mode 100644 index 000000000..fbd923d91 --- /dev/null +++ b/python/sglang/srt/entrypoints/engine.py @@ -0,0 +1,872 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +""" +The entry point of inference server. (SRT = SGLang Runtime) + +This file implements python APIs for the inference engine. +""" + +import asyncio +import atexit +import dataclasses +import logging +import multiprocessing as mp +import os +import random +import signal +import threading +import time +from typing import AsyncIterator, Dict, Iterator, List, Optional, Tuple, Union + +import zmq +import zmq.asyncio +from PIL.Image import Image + +# Fix a bug of Python threading +setattr(threading, "_register_atexit", lambda *args, **kwargs: None) + +import torch +import uvloop + +from sglang.srt.entrypoints.EngineBase import EngineBase +from sglang.srt.managers.data_parallel_controller import ( + run_data_parallel_controller_process, +) +from sglang.srt.managers.detokenizer_manager import run_detokenizer_process +from sglang.srt.managers.io_struct import ( + EmbeddingReqInput, + GenerateReqInput, + GetWeightsByNameReqInput, + InitWeightsUpdateGroupReqInput, + LoadLoRAAdapterReqInput, + MultimodalDataInputFormat, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, + RpcReqInput, + RpcReqOutput, + UnloadLoRAAdapterReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromTensorReqInput, +) +from sglang.srt.managers.multi_tokenizer_mixin import MultiTokenizerRouter +from sglang.srt.managers.scheduler import run_scheduler_process +from sglang.srt.managers.template_manager import TemplateManager +from sglang.srt.managers.tokenizer_manager import TokenizerManager +from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter +from sglang.srt.utils import ( + MultiprocessingSerializer, + assert_pkg_version, + configure_logger, + get_bool_env_var, + get_zmq_socket, + is_cuda, + kill_process_tree, + launch_dummy_health_check_server, + prepare_model_and_tokenizer, + set_prometheus_multiproc_dir, + set_ulimit, +) +from sglang.version import __version__ + +logger = logging.getLogger(__name__) +asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + +_is_cuda = is_cuda() + + +class Engine(EngineBase): + """ + The entry point to the inference engine. + + - The engine consists of three components: + 1. TokenizerManager: Tokenizes the requests and sends them to the scheduler. + 2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager. + 3. DetokenizerManager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager. + + Note: + 1. The HTTP server, Engine, and TokenizerManager all run in the main process. + 2. Inter-process communication (IPC) is handled via the ZMQ library, with each process using a different port. + """ + + def __init__(self, **kwargs): + """ + The arguments of this function is the same as `sglang/srt/server_args.py::ServerArgs`. + Please refer to `ServerArgs` for the documentation. + """ + if "server_args" in kwargs: + # Directly load server_args + server_args = kwargs["server_args"] + else: + # Construct server_args from kwargs + if "log_level" not in kwargs: + # Do not print logs by default + kwargs["log_level"] = "error" + server_args = ServerArgs(**kwargs) + + # Shutdown the subprocesses automatically when the program exits + atexit.register(self.shutdown) + + # Allocate ports for inter-process communications + self.port_args = PortArgs.init_new(server_args) + logger.info(f"{server_args=}") + + # Launch subprocesses + tokenizer_manager, template_manager, scheduler_info = _launch_subprocesses( + server_args=server_args, + port_args=self.port_args, + ) + self.server_args = server_args + self.tokenizer_manager = tokenizer_manager + self.template_manager = template_manager + self.scheduler_info = scheduler_info + + context = zmq.Context(2) + self.send_to_rpc = get_zmq_socket( + context, zmq.DEALER, self.port_args.rpc_ipc_name, True + ) + + def generate( + self, + # The input prompt. It can be a single prompt or a batch of prompts. + prompt: Optional[Union[List[str], str]] = None, + sampling_params: Optional[Union[List[Dict], Dict]] = None, + # The token ids for text; one can either specify text or input_ids. + input_ids: Optional[Union[List[List[int]], List[int]]] = None, + # The image input. It can be an image instance, file name, URL, or base64 encoded string. + # Can be formatted as: + # - Single image for a single request + # - List of images (one per request in a batch) + # - List of lists of images (multiple images per request) + # See also python/sglang/srt/utils.py:load_image for more details. + image_data: Optional[MultimodalDataInputFormat] = None, + audio_data: Optional[MultimodalDataInputFormat] = None, + video_data: Optional[MultimodalDataInputFormat] = None, + return_logprob: Optional[Union[List[bool], bool]] = False, + logprob_start_len: Optional[Union[List[int], int]] = None, + top_logprobs_num: Optional[Union[List[int], int]] = None, + token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None, + lora_path: Optional[List[Optional[str]]] = None, + custom_logit_processor: Optional[Union[List[str], str]] = None, + return_hidden_states: bool = False, + stream: bool = False, + bootstrap_host: Optional[Union[List[str], str]] = None, + bootstrap_port: Optional[Union[List[int], int]] = None, + bootstrap_room: Optional[Union[List[int], int]] = None, + data_parallel_rank: Optional[int] = None, + ) -> Union[Dict, Iterator[Dict]]: + """ + The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`. + Please refer to `GenerateReqInput` for the documentation. + """ + if self.server_args.enable_dp_attention: + if data_parallel_rank is None: + logger.debug("data_parallel_rank not provided, using default dispatch") + elif data_parallel_rank < 0: + raise ValueError("data_parallel_rank must be non-negative") + elif data_parallel_rank >= self.server_args.dp_size: + raise ValueError( + f"data_parallel_rank must be less than dp_size: {self.server_args.dp_size}" + ) + + obj = GenerateReqInput( + text=prompt, + input_ids=input_ids, + sampling_params=sampling_params, + image_data=image_data, + audio_data=audio_data, + video_data=video_data, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + token_ids_logprob=token_ids_logprob, + lora_path=lora_path, + custom_logit_processor=custom_logit_processor, + return_hidden_states=return_hidden_states, + stream=stream, + bootstrap_host=bootstrap_host, + bootstrap_port=bootstrap_port, + bootstrap_room=bootstrap_room, + data_parallel_rank=data_parallel_rank, + ) + loop = asyncio.get_event_loop() + generator = self.tokenizer_manager.generate_request(obj, None) + + if stream: + + def generator_wrapper(): + while True: + try: + chunk = loop.run_until_complete(generator.__anext__()) + yield chunk + except StopAsyncIteration: + break + + return generator_wrapper() + else: + ret = loop.run_until_complete(generator.__anext__()) + return ret + + async def async_generate( + self, + # The input prompt. It can be a single prompt or a batch of prompts. + prompt: Optional[Union[List[str], str]] = None, + sampling_params: Optional[Union[List[Dict], Dict]] = None, + # The token ids for text; one can either specify text or input_ids. + input_ids: Optional[Union[List[List[int]], List[int]]] = None, + # The image input. It can be an image instance, file name, URL, or base64 encoded string. + # Can be formatted as: + # - Single image for a single request + # - List of images (one per request in a batch) + # - List of lists of images (multiple images per request) + # See also python/sglang/srt/utils.py:load_image for more details. + image_data: Optional[MultimodalDataInputFormat] = None, + audio_data: Optional[MultimodalDataInputFormat] = None, + video_data: Optional[MultimodalDataInputFormat] = None, + return_logprob: Optional[Union[List[bool], bool]] = False, + logprob_start_len: Optional[Union[List[int], int]] = None, + top_logprobs_num: Optional[Union[List[int], int]] = None, + token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None, + lora_path: Optional[List[Optional[str]]] = None, + custom_logit_processor: Optional[Union[List[str], str]] = None, + return_hidden_states: bool = False, + stream: bool = False, + bootstrap_host: Optional[Union[List[str], str]] = None, + bootstrap_port: Optional[Union[List[int], int]] = None, + bootstrap_room: Optional[Union[List[int], int]] = None, + data_parallel_rank: Optional[int] = None, + ) -> Union[Dict, AsyncIterator[Dict]]: + """ + The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`. + Please refer to `GenerateReqInput` for the documentation. + """ + + if self.server_args.enable_dp_attention: + if data_parallel_rank is None: + logger.debug("data_parallel_rank not provided, using default dispatch") + elif data_parallel_rank < 0: + raise ValueError("data_parallel_rank must be non-negative") + elif data_parallel_rank >= self.server_args.dp_size: + raise ValueError( + f"data_parallel_rank must be in range [0, {self.server_args.dp_size-1}]" + ) + + logger.debug(f"data_parallel_rank: {data_parallel_rank}") + obj = GenerateReqInput( + text=prompt, + input_ids=input_ids, + sampling_params=sampling_params, + image_data=image_data, + audio_data=audio_data, + video_data=video_data, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + token_ids_logprob=token_ids_logprob, + lora_path=lora_path, + return_hidden_states=return_hidden_states, + stream=stream, + custom_logit_processor=custom_logit_processor, + bootstrap_host=bootstrap_host, + bootstrap_port=bootstrap_port, + bootstrap_room=bootstrap_room, + data_parallel_rank=data_parallel_rank, + ) + generator = self.tokenizer_manager.generate_request(obj, None) + + if stream is True: + return generator + else: + return await generator.__anext__() + + def encode( + self, + prompt: Union[str, List[str], List[Dict], List[List[Dict]]], + image_data: Optional[MultimodalDataInputFormat] = None, + audio_data: Optional[MultimodalDataInputFormat] = None, + video_data: Optional[MultimodalDataInputFormat] = None, + ) -> Dict: + """ + The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`. + Please refer to `EmbeddingReqInput` for the documentation. + """ + obj = EmbeddingReqInput( + text=prompt, + image_data=image_data, + audio_data=audio_data, + video_data=video_data, + ) + loop = asyncio.get_event_loop() + generator = self.tokenizer_manager.generate_request(obj, None) + ret = loop.run_until_complete(generator.__anext__()) + return ret + + async def async_encode( + self, + prompt: Union[str, List[str], List[Dict], List[List[Dict]]], + image_data: Optional[MultimodalDataInputFormat] = None, + audio_data: Optional[MultimodalDataInputFormat] = None, + video_data: Optional[MultimodalDataInputFormat] = None, + ) -> Dict: + """ + Asynchronous version of encode method. + + The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`. + Please refer to `EmbeddingReqInput` for the documentation. + """ + obj = EmbeddingReqInput( + text=prompt, + image_data=image_data, + audio_data=audio_data, + video_data=video_data, + ) + generator = self.tokenizer_manager.generate_request(obj, None) + return await generator.__anext__() + + def rerank( + self, + prompt: Union[List[List[str]]], + ) -> Dict: + """ + The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`. + Please refer to `EmbeddingReqInput` for the documentation. + """ + obj = EmbeddingReqInput(text=prompt, is_cross_encoder_request=True) + loop = asyncio.get_event_loop() + generator = self.tokenizer_manager.generate_request(obj, None) + ret = loop.run_until_complete(generator.__anext__()) + return ret + + def shutdown(self): + """Shutdown the engine""" + kill_process_tree(os.getpid(), include_parent=False) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.shutdown() + return False + + def flush_cache(self): + loop = asyncio.get_event_loop() + return loop.run_until_complete(self.tokenizer_manager.flush_cache()) + + def start_profile(self): + loop = asyncio.get_event_loop() + loop.run_until_complete(self.tokenizer_manager.start_profile()) + + def stop_profile(self): + loop = asyncio.get_event_loop() + loop.run_until_complete(self.tokenizer_manager.stop_profile()) + + def start_expert_distribution_record(self): + loop = asyncio.get_event_loop() + loop.run_until_complete( + self.tokenizer_manager.start_expert_distribution_record() + ) + + def stop_expert_distribution_record(self): + loop = asyncio.get_event_loop() + loop.run_until_complete( + self.tokenizer_manager.stop_expert_distribution_record() + ) + + def dump_expert_distribution_record(self): + loop = asyncio.get_event_loop() + loop.run_until_complete( + self.tokenizer_manager.dump_expert_distribution_record() + ) + + def get_server_info(self): + loop = asyncio.get_event_loop() + internal_states = loop.run_until_complete( + self.tokenizer_manager.get_internal_state() + ) + return { + **dataclasses.asdict(self.tokenizer_manager.server_args), + **self.scheduler_info, + "internal_states": internal_states, + "version": __version__, + } + + def init_weights_update_group( + self, + master_address: str, + master_port: int, + rank_offset: int, + world_size: int, + group_name: str, + backend: str = "nccl", + ): + """Initialize parameter update group.""" + obj = InitWeightsUpdateGroupReqInput( + master_address=master_address, + master_port=master_port, + rank_offset=rank_offset, + world_size=world_size, + group_name=group_name, + backend=backend, + ) + loop = asyncio.get_event_loop() + return loop.run_until_complete( + self.tokenizer_manager.init_weights_update_group(obj, None) + ) + + def update_weights_from_distributed( + self, + names: list[str], + dtypes: list[str], + shapes: list[list[int]], + group_name: str = "weight_update_group", + flush_cache: bool = True, + ): + """Update weights from distributed source.""" + obj = UpdateWeightsFromDistributedReqInput( + names=names, + dtypes=dtypes, + shapes=shapes, + group_name=group_name, + flush_cache=flush_cache, + ) + loop = asyncio.get_event_loop() + return loop.run_until_complete( + self.tokenizer_manager.update_weights_from_distributed(obj, None) + ) + + def update_weights_from_tensor( + self, + named_tensors: List[Tuple[str, torch.Tensor]], + load_format: Optional[str] = None, + flush_cache: bool = True, + ): + """Update weights from distributed source. If there are going to be more updates, set `flush_cache` to be false + to avoid duplicated cache cleaning operation.""" + if load_format == "flattened_bucket": + serialized_named_tensors = named_tensors + else: + serialized_named_tensors = [ + MultiprocessingSerializer.serialize(named_tensors) + for _ in range(self.server_args.tp_size) + ] + obj = UpdateWeightsFromTensorReqInput( + serialized_named_tensors=serialized_named_tensors, + load_format=load_format, + flush_cache=flush_cache, + ) + loop = asyncio.get_event_loop() + + return loop.run_until_complete( + self.tokenizer_manager.update_weights_from_tensor(obj, None) + ) + + def update_weights_from_disk( + self, + model_path: str, + load_format: Optional[str] = None, + ): + """Update the weights from disk inplace without re-launching the engine. + + This method allows updating the model weights from disk without restarting + the engine. It can be used to load a different model or update weights with + new training. + """ + obj = UpdateWeightFromDiskReqInput( + model_path=model_path, + load_format=load_format, + ) + + loop = asyncio.get_event_loop() + return loop.run_until_complete( + self.tokenizer_manager.update_weights_from_disk(obj, None) + ) + + def get_weights_by_name(self, name: str, truncate_size: int = 100): + """Get weights by parameter name.""" + obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size) + loop = asyncio.get_event_loop() + return loop.run_until_complete( + self.tokenizer_manager.get_weights_by_name(obj, None) + ) + + def load_lora_adapter(self, lora_name: str, lora_path: str, pinned: bool = False): + """Load a new LoRA adapter without re-launching the engine.""" + + obj = LoadLoRAAdapterReqInput( + lora_name=lora_name, + lora_path=lora_path, + pinned=pinned, + ) + + loop = asyncio.get_event_loop() + return loop.run_until_complete( + self.tokenizer_manager.load_lora_adapter(obj, None) + ) + + def unload_lora_adapter(self, lora_name: str): + """Unload a LoRA adapter without re-launching the engine.""" + + obj = UnloadLoRAAdapterReqInput(lora_name=lora_name) + + loop = asyncio.get_event_loop() + return loop.run_until_complete( + self.tokenizer_manager.unload_lora_adapter(obj, None) + ) + + def release_memory_occupation(self, tags: Optional[List[str]] = None): + obj = ReleaseMemoryOccupationReqInput(tags=tags) + loop = asyncio.get_event_loop() + return loop.run_until_complete( + self.tokenizer_manager.release_memory_occupation(obj, None) + ) + + def resume_memory_occupation(self, tags: Optional[List[str]] = None): + obj = ResumeMemoryOccupationReqInput(tags=tags) + loop = asyncio.get_event_loop() + return loop.run_until_complete( + self.tokenizer_manager.resume_memory_occupation(obj, None) + ) + + def freeze_gc(self): + """ + To maintain a high performance server with low latency, we want to reduce the + stalls caused by the garbage collector scanning through a large number of objects. + + It is usually helpful to start the server and warm it up with real requests to + initialize many of the long-lived objects that do not need to be garbage collected. + + After sufficient warmup, we can call this function to freeze the garbage collector + so that all objects created before this point are considered out of scope for garbage + collection. + """ + + loop = asyncio.get_event_loop() + loop.run_until_complete(self.tokenizer_manager.freeze_gc()) + + """ + Execute an RPC call on all scheduler processes. + """ + + def collective_rpc(self, method: str, **kwargs): + obj = RpcReqInput(method=method, parameters=kwargs) + self.send_to_rpc.send_pyobj(obj) + recv_req = self.send_to_rpc.recv_pyobj(zmq.BLOCKY) + assert isinstance(recv_req, RpcReqOutput) + assert recv_req.success, recv_req.message + + def save_remote_model(self, **kwargs): + self.collective_rpc("save_remote_model", **kwargs) + + def save_sharded_model(self, **kwargs): + self.collective_rpc("save_sharded_model", **kwargs) + + def score( + self, + query: Optional[Union[str, List[int]]] = None, + items: Optional[Union[str, List[str], List[List[int]]]] = None, + label_token_ids: Optional[List[int]] = None, + apply_softmax: bool = False, + item_first: bool = False, + ) -> List[List[float]]: + """ + Score the probability of specified token IDs appearing after the given (query + item) pair. For example: + query = "<|user|>Is the following city the capital of France? " + items = ["Paris <|assistant|>", "London <|assistant|>", "Berlin <|assistant|>"] + label_token_ids = [2332, 1223] # Token IDs for "Yes" and "No" + item_first = False + + This would pass the following prompts to the model: + "<|user|>Is the following city the capital of France? Paris <|assistant|>" + "<|user|>Is the following city the capital of France? London <|assistant|>" + "<|user|>Is the following city the capital of France? Berlin <|assistant|>" + The api would then return the probabilities of the model producing "Yes" and "No" as the next token. + The output would look like: + [[0.9, 0.1], [0.2, 0.8], [0.1, 0.9]] + + + Args: + query: The query text or pre-tokenized query token IDs. Must be provided. + items: The item text(s) or pre-tokenized item token IDs. Must be provided. + label_token_ids: List of token IDs to compute probabilities for. If None, no token probabilities will be computed. + apply_softmax: Whether to normalize probabilities using softmax. + item_first: If True, prepend items to query. Otherwise append items to query. + + Returns: + List of dictionaries mapping token IDs to their probabilities for each item. + Each dictionary in the list corresponds to one item input. + + Raises: + ValueError: If query is not provided, or if items is not provided, + or if token IDs are out of vocabulary, or if logprobs are not available for the specified tokens. + """ + loop = asyncio.get_event_loop() + return loop.run_until_complete( + self.tokenizer_manager.score_request( + query=query, + items=items, + label_token_ids=label_token_ids, + apply_softmax=apply_softmax, + item_first=item_first, + request=None, + ) + ) + + async def async_score( + self, + query: Optional[Union[str, List[int]]] = None, + items: Optional[Union[str, List[str], List[List[int]]]] = None, + label_token_ids: Optional[List[int]] = None, + apply_softmax: bool = False, + item_first: bool = False, + ) -> List[List[float]]: + """ + Asynchronous version of score method. + + See score() for detailed documentation. + """ + return await self.tokenizer_manager.score_request( + query=query, + items=items, + label_token_ids=label_token_ids, + apply_softmax=apply_softmax, + item_first=item_first, + request=None, + ) + + +def _set_envs_and_config(server_args: ServerArgs): + # Set global environments + os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" + os.environ["NCCL_CUMEM_ENABLE"] = str(int(server_args.enable_symm_mem)) + if not server_args.enable_symm_mem: + os.environ["NCCL_NVLS_ENABLE"] = str(int(server_args.enable_nccl_nvls)) + os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "4" + os.environ["CUDA_MODULE_LOADING"] = "AUTO" + # flashinfer uses this environment variable for various kernels from MoE to quant kernels + if os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0": + os.environ["TRTLLM_ENABLE_PDL"] = "1" + + # Can also be passed as argument + os.environ["SGLANG_RUN_ID"] = ( + f"sglang-run-{time.time()}-{random.randint(0, 100000000)}" + ) + + # Set prometheus env vars + if server_args.enable_metrics: + set_prometheus_multiproc_dir() + + # Set ulimit + set_ulimit() + + # Check flashinfer version + if server_args.attention_backend == "flashinfer": + assert_pkg_version( + "flashinfer_python", + "0.3.1", + "Please uninstall the old version and " + "reinstall the latest version by following the instructions " + "at https://docs.flashinfer.ai/installation.html.", + ) + if _is_cuda and not get_bool_env_var("SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK"): + assert_pkg_version( + "sgl-kernel", + "0.3.9.post2", + "Please reinstall the latest version with `pip install sgl-kernel --force-reinstall`", + ) + + if True: # Keep this check for internal code compatibility + # Register the signal handler. + # The child processes will send SIGQUIT to this process when any error happens + # This process then clean up the whole process tree + # Note: This sigquit handler is used in the launch phase, and may be replaced by + # the running_phase_sigquit_handler in the tokenizer manager after the grpc server is launched. + def launch_phase_sigquit_handler(signum, frame): + logger.error( + "Received sigquit from a child process. It usually means the child failed." + ) + kill_process_tree(os.getpid()) + + signal.signal(signal.SIGQUIT, launch_phase_sigquit_handler) + + # Set mp start method + mp.set_start_method("spawn", force=True) + + +def _init_tokenizer_manager( + server_args: ServerArgs, port_args: PortArgs +) -> TokenizerManager: + # Launch tokenizer process + tokenizer_manager = TokenizerManager(server_args, port_args) + + # Initialize templates + template_manager = TemplateManager() + template_manager.initialize_templates( + tokenizer_manager=tokenizer_manager, + model_path=server_args.model_path, + chat_template=server_args.chat_template, + completion_template=server_args.completion_template, + ) + + return tokenizer_manager, template_manager + + +def _launch_subprocesses( + server_args: ServerArgs, port_args: Optional[PortArgs] = None +) -> Tuple[TokenizerManager, TemplateManager, Dict]: + """ + Launch the TokenizerManager in the main process, the Scheduler in a subprocess, and the DetokenizerManager in another subprocess. + """ + # Configure global environment + configure_logger(server_args) + server_args.check_server_args() + _set_envs_and_config(server_args) + + # Allocate ports for inter-process communications + if port_args is None: + port_args = PortArgs.init_new(server_args) + logger.info(f"{server_args=}") + + # If using model from www.modelscope.cn, first download the model. + server_args.model_path, server_args.tokenizer_path = prepare_model_and_tokenizer( + server_args.model_path, server_args.tokenizer_path + ) + + scheduler_procs = [] + if server_args.dp_size == 1: + memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=server_args.enable_memory_saver + ) + scheduler_pipe_readers = [] + + nnodes_per_tp_group = max(server_args.nnodes // server_args.pp_size, 1) + tp_size_per_node = server_args.tp_size // nnodes_per_tp_group + tp_rank_range = range( + tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group), + tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1), + ) + + pp_size_per_node = max(server_args.pp_size // server_args.nnodes, 1) + pp_rank_range = range( + pp_size_per_node * (server_args.node_rank // nnodes_per_tp_group), + pp_size_per_node * (server_args.node_rank // nnodes_per_tp_group + 1), + ) + + for pp_rank in pp_rank_range: + for tp_rank in tp_rank_range: + reader, writer = mp.Pipe(duplex=False) + gpu_id = ( + server_args.base_gpu_id + + ((pp_rank % pp_size_per_node) * tp_size_per_node) + + (tp_rank % tp_size_per_node) * server_args.gpu_id_step + ) + moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size) + proc = mp.Process( + target=run_scheduler_process, + args=( + server_args, + port_args, + gpu_id, + tp_rank, + moe_ep_rank, + pp_rank, + None, + writer, + None, + ), + ) + + with memory_saver_adapter.configure_subprocess(): + proc.start() + scheduler_procs.append(proc) + scheduler_pipe_readers.append(reader) + else: + # Launch the data parallel controller + reader, writer = mp.Pipe(duplex=False) + scheduler_pipe_readers = [reader] + proc = mp.Process( + target=run_data_parallel_controller_process, + args=(server_args, port_args, writer), + ) + proc.start() + scheduler_procs.append(proc) + + if server_args.node_rank >= 1: + # In multi-node cases, non-zero rank nodes do not need to run tokenizer or detokenizer, + # so they can just wait here. + + for reader in scheduler_pipe_readers: + data = reader.recv() + assert data["status"] == "ready" + + if os.getenv("SGLANG_BLOCK_NONZERO_RANK_CHILDREN") == "0": + # When using `Engine` as a Python API, we don't want to block here. + return None, None, None + + launch_dummy_health_check_server( + server_args.host, server_args.port, server_args.enable_metrics + ) + + for proc in scheduler_procs: + proc.join() + logger.error( + f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}" + ) + return None, None, None + + # Launch detokenizer process + detoken_proc = mp.Process( + target=run_detokenizer_process, + args=( + server_args, + port_args, + ), + ) + detoken_proc.start() + + # Init tokenizer manager first, as the bootstrap server is initialized here + if server_args.tokenizer_worker_num > 1: + # Launch multi-tokenizer router + tokenizer_manager = MultiTokenizerRouter(server_args, port_args) + template_manager = None + else: + tokenizer_manager, template_manager = _init_tokenizer_manager( + server_args, port_args + ) + + # Wait for the model to finish loading + scheduler_infos = [] + for i in range(len(scheduler_pipe_readers)): + try: + data = scheduler_pipe_readers[i].recv() + except EOFError: + logger.error( + f"Rank {i} scheduler is dead. Please check if there are relevant logs." + ) + scheduler_procs[i].join() + logger.error(f"Exit code: {scheduler_procs[i].exitcode}") + raise + + if data["status"] != "ready": + raise RuntimeError( + "Initialization failed. Please see the error messages above." + ) + scheduler_infos.append(data) + + # Assume all schedulers have the same scheduler_info + scheduler_info = scheduler_infos[0] + + tokenizer_manager.max_req_input_len = scheduler_info["max_req_input_len"] + + return tokenizer_manager, template_manager, scheduler_info diff --git a/python/sglang/srt/entrypoints/harmony_utils.py b/python/sglang/srt/entrypoints/harmony_utils.py new file mode 100644 index 000000000..5ebb653b3 --- /dev/null +++ b/python/sglang/srt/entrypoints/harmony_utils.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from vLLM: https://github.com/vllm-project/vllm/blob/1b9902806915040ac9b3029f2ab7522ec505afc3/vllm/entrypoints/harmony_utils.py +# Slight differences in processing chat messages +import datetime +import json +from collections.abc import Iterable +from typing import Literal, Optional, Union + +from openai.types.responses import ( + ResponseOutputItem, + ResponseOutputMessage, + ResponseOutputText, + ResponseReasoningItem, +) +from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall +from openai.types.responses.response_function_web_search import ( + ActionFind, + ActionOpenPage, + ActionSearch, + ResponseFunctionWebSearch, +) +from openai.types.responses.response_reasoning_item import ( + Content as ResponseReasoningTextContent, +) +from openai.types.responses.tool import Tool +from openai_harmony import ( + Author, + Conversation, + DeveloperContent, + HarmonyEncodingName, + Message, + ReasoningEffort, + Role, + StreamableParser, + SystemContent, + TextContent, + ToolDescription, + load_harmony_encoding, +) + +from sglang.srt.entrypoints.openai.protocol import ResponseInputOutputItem +from sglang.srt.utils import random_uuid + +REASONING_EFFORT = { + "high": ReasoningEffort.HIGH, + "medium": ReasoningEffort.MEDIUM, + "low": ReasoningEffort.LOW, +} + +_harmony_encoding = None + + +def get_encoding(): + global _harmony_encoding + if _harmony_encoding is None: + _harmony_encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS) + return _harmony_encoding + + +def get_system_message( + model_identity: Optional[str] = None, + reasoning_effort: Optional[Literal["high", "medium", "low"]] = None, + start_date: Optional[str] = None, + browser_description: Optional[str] = None, + python_description: Optional[str] = None, +) -> Message: + sys_msg_content = SystemContent.new() + if model_identity is not None: + sys_msg_content = sys_msg_content.with_model_identity(model_identity) + if reasoning_effort is not None: + sys_msg_content = sys_msg_content.with_reasoning_effort( + REASONING_EFFORT[reasoning_effort] + ) + if start_date is None: + start_date = datetime.datetime.now().strftime("%Y-%m-%d") + sys_msg_content = sys_msg_content.with_conversation_start_date(start_date) + if browser_description is not None: + sys_msg_content = sys_msg_content.with_tools(browser_description) + if python_description is not None: + sys_msg_content = sys_msg_content.with_tools(python_description) + sys_msg = Message.from_role_and_content(Role.SYSTEM, sys_msg_content) + return sys_msg + + +def get_developer_message( + instructions: Optional[str] = None, tools: Optional[list[Tool]] = None +) -> Message: + dev_msg_content = DeveloperContent.new() + if instructions is not None: + dev_msg_content = dev_msg_content.with_instructions(instructions) + if tools is not None: + function_tools = [] + for tool in tools: + if tool.type in ("web_search_preview", "code_interpreter"): + # These are built-in tools that are added to the system message. + pass + elif tool.type == "function": + function_tools.append(tool) + else: + raise ValueError(f"tool type {tool.type} not supported") + if function_tools: + function_tool_descriptions = [ + ToolDescription.new( + name=tool.name, + description=tool.description, + parameters=tool.parameters, + ) + for tool in function_tools + ] + dev_msg_content = dev_msg_content.with_function_tools( + function_tool_descriptions + ) + dev_msg = Message.from_role_and_content(Role.DEVELOPER, dev_msg_content) + return dev_msg + + +def get_user_message(content: str) -> Message: + return Message.from_role_and_content(Role.USER, content) + + +def parse_response_input( + response_msg: ResponseInputOutputItem, + prev_responses: list[Union[ResponseOutputItem, ResponseReasoningItem]], +) -> Message: + if not isinstance(response_msg, dict): + response_msg = response_msg.model_dump() + if "type" not in response_msg or response_msg["type"] == "message": + role = response_msg["role"] + content = response_msg["content"] + if role == "system": + # User is trying to set a system message. Change it to: + # <|start|>developer<|message|># Instructions + # {instructions}<|end|> + role = "developer" + text_prefix = "Instructions:\n" + else: + text_prefix = "" + if isinstance(content, str): + msg = Message.from_role_and_content(role, text_prefix + content) + else: + contents = [TextContent(text=text_prefix + c["text"]) for c in content] + msg = Message.from_role_and_contents(role, contents) + elif response_msg["type"] == "function_call_output": + call_id = response_msg["call_id"] + call_response: Optional[ResponseFunctionToolCall] = None + for prev_response in reversed(prev_responses): + if ( + isinstance(prev_response, ResponseFunctionToolCall) + and prev_response.call_id == call_id + ): + call_response = prev_response + break + if call_response is None: + raise ValueError(f"No call message found for {call_id}") + msg = Message.from_author_and_content( + Author.new(Role.TOOL, f"functions.{call_response.name}"), + response_msg["output"], + ) + elif response_msg["type"] == "reasoning": + content = response_msg["content"] + assert len(content) == 1 + msg = Message.from_role_and_content(Role.ASSISTANT, content[0]["text"]) + elif response_msg["type"] == "function_call": + msg = Message.from_role_and_content(Role.ASSISTANT, response_msg["arguments"]) + msg = msg.with_channel("commentary") + msg = msg.with_recipient(f"functions.{response_msg['name']}") + msg = msg.with_content_type("json") + else: + raise ValueError(f"Unknown input type: {response_msg['type']}") + return msg + + +def parse_response_output(output: ResponseOutputItem) -> Message: + if isinstance(output, ResponseOutputMessage): + role = output.role + contents = [TextContent(text=c.text) for c in output.content] + msg = Message.from_role_and_contents(role, contents) + return msg + elif isinstance(output, ResponseFunctionToolCall): + msg = Message.from_role_and_content(Role.ASSISTANT, output.arguments) + msg = msg.with_channel("commentary") + msg = msg.with_recipient(output.name) + msg = msg.with_content_type("json") + return msg + else: + raise ValueError(f"Unknown output type: {type(output)}") + + +def parse_chat_input(chat_msg) -> Message: + role = chat_msg.role + content = chat_msg.content + if isinstance(content, str): + contents = [TextContent(text=content)] + else: + # TODO: Support refusal. + contents = [TextContent(text=c.text) for c in content] + msg = Message.from_role_and_contents(role, contents) + return msg + + +def render_for_completion(messages: list[Message]) -> list[int]: + conversation = Conversation.from_messages(messages) + token_ids = get_encoding().render_conversation_for_completion( + conversation, Role.ASSISTANT + ) + return token_ids + + +def get_stop_tokens_for_assistant_actions() -> list[int]: + return get_encoding().stop_tokens_for_assistant_actions() + + +def get_streamable_parser_for_assistant() -> StreamableParser: + return StreamableParser(get_encoding(), role=Role.ASSISTANT) + + +def parse_output_message(message: Message): + if message.author.role != "assistant": + # This is a message from a tool to the assistant (e.g., search result). + # Don't include it in the final output for now. This aligns with + # OpenAI's behavior on models like o4-mini. + return [] + + output_items = [] + recipient = message.recipient + if recipient is not None and recipient.startswith("browser."): + if len(message.content) != 1: + raise ValueError("Invalid number of contents in browser message") + content = message.content[0] + browser_call = json.loads(content.text) + # TODO: translate to url properly! + if recipient == "browser.search": + action = ActionSearch( + query=f"cursor:{browser_call.get('query', '')}", type="search" + ) + elif recipient == "browser.open": + action = ActionOpenPage( + url=f"cursor:{browser_call.get('url', '')}", type="open_page" + ) + elif recipient == "browser.find": + action = ActionFind( + pattern=browser_call["pattern"], + url=f"cursor:{browser_call.get('url', '')}", + type="find", + ) + else: + raise ValueError(f"Unknown browser action: {recipient}") + web_search_item = ResponseFunctionWebSearch( + id=f"ws_{random_uuid()}", + action=action, + status="completed", + type="web_search_call", + ) + output_items.append(web_search_item) + elif message.channel == "analysis": + for content in message.content: + reasoning_item = ResponseReasoningItem( + id=f"rs_{random_uuid()}", + type="reasoning", + summary=[], + content=[ + ResponseReasoningTextContent( + text=content.text, type="reasoning_text" + ) + ], + status=None, + ) + output_items.append(reasoning_item) + elif message.channel == "commentary": + if message.recipient.startswith("functions."): + function_name = message.recipient.split(".")[-1] + for content in message.content: + random_id = random_uuid() + response_item = ResponseFunctionToolCall( + arguments=content.text, + call_id=f"call_{random_id}", + type="function_call", + name=function_name, + id=f"ft_{random_id}", + ) + output_items.append(response_item) + elif message.recipient.startswith("python") or message.recipient.startswith( + "browser" + ): + for content in message.content: + reasoning_item = ResponseReasoningItem( + id=f"rs_{random_uuid()}", + type="reasoning", + summary=[], + content=[ + ResponseReasoningTextContent( + text=content.text, type="reasoning_text" + ) + ], + status=None, + ) + output_items.append(reasoning_item) + else: + raise ValueError(f"Unknown recipient: {message.recipient}") + elif message.channel == "final": + contents = [] + for content in message.content: + output_text = ResponseOutputText( + text=content.text, + annotations=[], # TODO + type="output_text", + logprobs=None, # TODO + ) + contents.append(output_text) + text_item = ResponseOutputMessage( + id=f"msg_{random_uuid()}", + content=contents, + role=message.author.role, + status="completed", + type="message", + ) + output_items.append(text_item) + else: + raise ValueError(f"Unknown channel: {message.channel}") + return output_items + + +def parse_remaining_state(parser: StreamableParser): + if not parser.current_content: + return [] + if parser.current_role != Role.ASSISTANT: + return [] + current_recipient = parser.current_recipient + if current_recipient is not None and current_recipient.startswith("browser."): + return [] + + if parser.current_channel == "analysis": + reasoning_item = ResponseReasoningItem( + id=f"rs_{random_uuid()}", + type="reasoning", + summary=[], + content=[ + ResponseReasoningTextContent( + text=parser.current_content, type="reasoning_text" + ) + ], + status=None, + ) + return [reasoning_item] + elif parser.current_channel == "final": + output_text = ResponseOutputText( + content=[ + ResponseReasoningTextContent( + text=parser.current_content, type="reasoning_text" + ) + ], + annotations=[], # TODO + type="output_text", + logprobs=None, # TODO + ) + text_item = ResponseOutputMessage( + id=f"msg_{random_uuid()}", + content=[output_text], + role="assistant", + status="completed", + type="message", + ) + return [text_item] + return [] + + +def parse_output_into_messages(token_ids: Iterable[int]): + parser = get_streamable_parser_for_assistant() + for token_id in token_ids: + parser.process(token_id) + return parser diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py new file mode 100644 index 000000000..9adac76ce --- /dev/null +++ b/python/sglang/srt/entrypoints/http_server.py @@ -0,0 +1,1398 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +""" +The entry point of inference server. (SRT = SGLang Runtime) + +This file implements HTTP APIs for the inference engine via fastapi. +""" + +import asyncio +import dataclasses +import json +import logging +import multiprocessing as multiprocessing +import os +import tempfile +import threading +import time +from http import HTTPStatus +from typing import Any, AsyncIterator, Callable, Dict, List, Optional + +import setproctitle + +# Fix a bug of Python threading +setattr(threading, "_register_atexit", lambda *args, **kwargs: None) + +from contextlib import asynccontextmanager +from typing import AsyncGenerator + +import numpy as np +import orjson +import requests +import uvicorn +import uvloop +from fastapi import Depends, FastAPI, HTTPException, Request, UploadFile +from fastapi.exceptions import RequestValidationError +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import ORJSONResponse, Response, StreamingResponse + +from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationMode +from sglang.srt.entrypoints.engine import _launch_subprocesses +from sglang.srt.entrypoints.openai.protocol import ( + ChatCompletionRequest, + CompletionRequest, + EmbeddingRequest, + ErrorResponse, + ModelCard, + ModelList, + ResponsesRequest, + ScoringRequest, + V1RerankReqInput, +) +from sglang.srt.entrypoints.openai.serving_chat import OpenAIServingChat +from sglang.srt.entrypoints.openai.serving_completions import OpenAIServingCompletion +from sglang.srt.entrypoints.openai.serving_embedding import OpenAIServingEmbedding +from sglang.srt.entrypoints.openai.serving_rerank import OpenAIServingRerank +from sglang.srt.entrypoints.openai.serving_score import OpenAIServingScore +from sglang.srt.function_call.function_call_parser import FunctionCallParser +from sglang.srt.managers.io_struct import ( + AbortReq, + CloseSessionReqInput, + ConfigureLoggingReq, + EmbeddingReqInput, + GenerateReqInput, + GetWeightsByNameReqInput, + InitWeightsUpdateGroupReqInput, + LoadLoRAAdapterReqInput, + OpenSessionReqInput, + ParseFunctionCallReq, + ProfileReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, + SeparateReasoningReqInput, + SetInternalStateReq, + SlowDownReqInput, + UnloadLoRAAdapterReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromTensorReqInput, + UpdateWeightVersionReqInput, + VertexGenerateReqInput, +) +from sglang.srt.managers.multi_tokenizer_mixin import ( + MultiTokenizerManager, + get_main_process_id, + monkey_patch_uvicorn_multiprocessing, + read_from_shared_memory, + write_data_for_multi_tokenizer, +) +from sglang.srt.managers.template_manager import TemplateManager +from sglang.srt.managers.tokenizer_manager import ServerStatus, TokenizerManager +from sglang.srt.metrics.func_timer import enable_func_timer +from sglang.srt.parser.reasoning_parser import ReasoningParser +from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.utils import ( + add_api_key_middleware, + add_prometheus_middleware, + delete_directory, + get_bool_env_var, + kill_process_tree, + set_uvicorn_logging_configs, +) +from sglang.srt.warmup import execute_warmups +from sglang.utils import get_exception_traceback +from sglang.version import __version__ + +logger = logging.getLogger(__name__) +asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + +HEALTH_CHECK_TIMEOUT = int(os.getenv("SGLANG_HEALTH_CHECK_TIMEOUT", 20)) + + +# Store global states +@dataclasses.dataclass +class _GlobalState: + tokenizer_manager: TokenizerManager + template_manager: TemplateManager + scheduler_info: Dict + + +_global_state: Optional[_GlobalState] = None + + +def set_global_state(global_state: _GlobalState): + global _global_state + _global_state = global_state + + +async def init_multi_tokenizer() -> ServerArgs: + """Read args information from shm and init tokenizer manager for current process""" + pid = os.getpid() + main_pid = get_main_process_id() + logger.info(f"current worker_id: {pid}, main processID: {main_pid}") + + # Read configuration from shared memory + port_args, server_args, scheduler_info = read_from_shared_memory( + f"multi_tokenizer_args_{main_pid}" + ) + server_args: ServerArgs + + # API key authentication is not supported in multi-tokenizer mode + assert ( + server_args.api_key is None + ), "API key is not supported in multi-tokenizer mode" + + port_args.tokenizer_ipc_name = ( + f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}" + ) + + # Launch multi-tokenizer manager process + tokenizer_manager = MultiTokenizerManager(server_args, port_args) + template_manager = TemplateManager() + template_manager.initialize_templates( + tokenizer_manager=tokenizer_manager, + model_path=server_args.model_path, + chat_template=server_args.chat_template, + completion_template=server_args.completion_template, + ) + # Register this tokenizer with the main tokenizer manager + await tokenizer_manager.register_to_main_tokenizer_manager() + + tokenizer_manager.max_req_input_len = scheduler_info["max_req_input_len"] + set_global_state( + _GlobalState( + tokenizer_manager=tokenizer_manager, + template_manager=template_manager, + scheduler_info=scheduler_info, + ) + ) + return server_args + + +@asynccontextmanager +async def lifespan(fast_api_app: FastAPI): + if not getattr(fast_api_app, "is_single_tokenizer_mode", False): + # Initialize multi-tokenizer support for worker processes + fast_api_app.server_args: ServerArgs = await init_multi_tokenizer() + + # only metrics middleware is supported in multi-tokenizer mode + worker_pid = os.getpid() + if fast_api_app.server_args.enable_metrics: + add_prometheus_middleware(app) + enable_func_timer() + + logger.info(f"Worker {worker_pid} added prometheus middleware") + fast_api_app.warmup_thread = threading.Thread( + target=_wait_and_warmup, + args=( + fast_api_app.server_args, + None, # pipe_finish_writer not needed in worker + None, # launch_callback not needed in worker + ), + ) + + # Initialize OpenAI serving handlers + fast_api_app.state.openai_serving_completion = OpenAIServingCompletion( + _global_state.tokenizer_manager, _global_state.template_manager + ) + fast_api_app.state.openai_serving_chat = OpenAIServingChat( + _global_state.tokenizer_manager, _global_state.template_manager + ) + fast_api_app.state.openai_serving_embedding = OpenAIServingEmbedding( + _global_state.tokenizer_manager, _global_state.template_manager + ) + fast_api_app.state.openai_serving_score = OpenAIServingScore( + _global_state.tokenizer_manager + ) + fast_api_app.state.openai_serving_rerank = OpenAIServingRerank( + _global_state.tokenizer_manager + ) + + server_args: ServerArgs = fast_api_app.server_args + + tool_server = None + if server_args.tool_server == "demo": + from sglang.srt.entrypoints.openai.tool_server import DemoToolServer + + tool_server = DemoToolServer() + elif server_args.tool_server: + from sglang.srt.entrypoints.openai.tool_server import MCPToolServer + + tool_server = MCPToolServer() + await tool_server.add_tool_server(server_args.tool_server) + + try: + from sglang.srt.entrypoints.openai.serving_responses import ( + OpenAIServingResponses, + ) + + fast_api_app.state.openai_serving_responses = OpenAIServingResponses( + _global_state.tokenizer_manager, + _global_state.template_manager, + enable_prompt_tokens_details=True, + enable_force_include_usage=True, + tool_server=tool_server, + ) + except Exception as e: + import traceback + + traceback.print_exc() + logger.warning(f"Can not initialize OpenAIServingResponses, error: {e}") + + if server_args.warmups is not None: + await execute_warmups( + server_args.disaggregation_mode, + server_args.warmups.split(","), + _global_state.tokenizer_manager, + ) + logger.info("Warmup ended") + + warmup_thread = getattr(fast_api_app, "warmup_thread", None) + if warmup_thread is not None: + warmup_thread.start() + + try: + yield + finally: + if server_args.tokenizer_worker_num > 1: + pid = os.getpid() + logger.info(f"uvicorn worker {pid} ending...") + warmup_thread.join() + logger.info(f"uvicorn worker {pid} ended.") + + +# Fast API +app = FastAPI( + lifespan=lifespan, + openapi_url=None if get_bool_env_var("DISABLE_OPENAPI_DOC") else "/openapi.json", +) +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.exception_handler(HTTPException) +async def validation_exception_handler(request: Request, exc: HTTPException): + """Enrich HTTP exception with status code and other details""" + error = ErrorResponse( + object="error", + message=exc.detail, + type=str(exc.status_code), + code=exc.status_code, + ) + return ORJSONResponse(content=error.model_dump(), status_code=exc.status_code) + + +# Custom exception handlers to change validation error status codes +@app.exception_handler(RequestValidationError) +async def validation_exception_handler(request: Request, exc: RequestValidationError): + """Override FastAPI's default 422 validation error with 400""" + exc_str = str(exc) + errors_str = str(exc.errors()) + + if errors_str and errors_str != exc_str: + message = f"{exc_str} {errors_str}" + else: + message = exc_str + + err = ErrorResponse( + message=message, + type=HTTPStatus.BAD_REQUEST.phrase, + code=HTTPStatus.BAD_REQUEST.value, + ) + + return ORJSONResponse( + status_code=400, + content=err.model_dump(), + ) + + +async def validate_json_request(raw_request: Request): + """Validate that the request content-type is application/json.""" + content_type = raw_request.headers.get("content-type", "").lower() + media_type = content_type.split(";", maxsplit=1)[0] + if media_type != "application/json": + raise RequestValidationError( + errors=[ + { + "loc": ["header", "content-type"], + "msg": "Unsupported Media Type: Only 'application/json' is allowed", + "type": "value_error", + } + ] + ) + + +##### Native API endpoints ##### + + +@app.get("/health") +@app.get("/health_generate") +async def health_generate(request: Request) -> Response: + """ + Check the health of the inference server by sending a special request to generate one token. + + If the server is running something, this request will be ignored, so it creates zero overhead. + If the server is not running anything, this request will be run, so we know whether the server is healthy. + """ + + if _global_state.tokenizer_manager.gracefully_exit: + logger.info("Health check request received during shutdown. Returning 503.") + return Response(status_code=503) + + if _global_state.tokenizer_manager.server_status == ServerStatus.Starting: + return Response(status_code=503) + + sampling_params = {"max_new_tokens": 1, "temperature": 0.0} + rid = f"HEALTH_CHECK_{time.time()}" + + if _global_state.tokenizer_manager.is_image_gen: + # Keep this branch for some internal use cases. + raise NotImplementedError("Image generation is not supported yet.") + elif _global_state.tokenizer_manager.is_generation: + gri = GenerateReqInput( + rid=rid, + input_ids=[0], + sampling_params=sampling_params, + log_metrics=False, + ) + if ( + _global_state.tokenizer_manager.server_args.disaggregation_mode + != DisaggregationMode.NULL + ): + gri.bootstrap_host = FAKE_BOOTSTRAP_HOST + gri.bootstrap_room = 0 + else: + gri = EmbeddingReqInput( + rid=rid, input_ids=[0], sampling_params=sampling_params, log_metrics=False + ) + + async def gen(): + async for _ in _global_state.tokenizer_manager.generate_request(gri, request): + break + + task = asyncio.create_task(gen()) + + # As long as we receive any response from the detokenizer/scheduler, we consider the server is healthy. + tic = time.time() + while time.time() < tic + HEALTH_CHECK_TIMEOUT: + await asyncio.sleep(1) + if _global_state.tokenizer_manager.last_receive_tstamp > tic: + task.cancel() + _global_state.tokenizer_manager.rid_to_state.pop(rid, None) + _global_state.tokenizer_manager.server_status = ServerStatus.Up + return Response(status_code=200) + + task.cancel() + tic_time = time.strftime("%H:%M:%S", time.localtime(tic)) + last_receive_time = time.strftime( + "%H:%M:%S", time.localtime(_global_state.tokenizer_manager.last_receive_tstamp) + ) + logger.error( + f"Health check failed. Server couldn't get a response from detokenizer for last " + f"{HEALTH_CHECK_TIMEOUT} seconds. tic start time: {tic_time}. " + f"last_heartbeat time: {last_receive_time}" + ) + _global_state.tokenizer_manager.rid_to_state.pop(rid, None) + _global_state.tokenizer_manager.server_status = ServerStatus.UnHealthy + return Response(status_code=503) + + +@app.get("/get_model_info") +async def get_model_info(): + """Get the model information.""" + result = { + "model_path": _global_state.tokenizer_manager.model_path, + "tokenizer_path": _global_state.tokenizer_manager.server_args.tokenizer_path, + "is_generation": _global_state.tokenizer_manager.is_generation, + "preferred_sampling_params": _global_state.tokenizer_manager.server_args.preferred_sampling_params, + "weight_version": _global_state.tokenizer_manager.server_args.weight_version, + } + return result + + +@app.get("/get_weight_version") +async def get_weight_version(): + """Get the current weight version.""" + return { + "weight_version": _global_state.tokenizer_manager.server_args.weight_version + } + + +@app.get("/get_server_info") +async def get_server_info(): + # Returns interna states per DP. + internal_states: List[Dict[Any, Any]] = ( + await _global_state.tokenizer_manager.get_internal_state() + ) + return { + **dataclasses.asdict(_global_state.tokenizer_manager.server_args), + **_global_state.scheduler_info, + "internal_states": internal_states, + "version": __version__, + } + + +@app.get("/get_load") +async def get_load(): + return await _global_state.tokenizer_manager.get_load() + + +# example usage: +# curl -s -X POST http://localhost:30000/set_internal_state -H "Content-Type: application/json" -d '{"server_args": {"max_micro_batch_size": 8}}' +@app.api_route("/set_internal_state", methods=["POST", "PUT"]) +async def set_internal_state(obj: SetInternalStateReq, request: Request): + res = await _global_state.tokenizer_manager.set_internal_state(obj) + return res + + +# fastapi implicitly converts json in the request to obj (dataclass) +@app.api_route("/generate", methods=["POST", "PUT"]) +async def generate_request(obj: GenerateReqInput, request: Request): + """Handle a generate request.""" + if obj.stream: + + async def stream_results() -> AsyncIterator[bytes]: + try: + async for out in _global_state.tokenizer_manager.generate_request( + obj, request + ): + yield b"data: " + orjson.dumps( + out, option=orjson.OPT_NON_STR_KEYS + ) + b"\n\n" + except ValueError as e: + out = {"error": {"message": str(e)}} + logger.error(f"[http_server] Error: {e}") + yield b"data: " + orjson.dumps( + out, option=orjson.OPT_NON_STR_KEYS + ) + b"\n\n" + yield b"data: [DONE]\n\n" + + return StreamingResponse( + stream_results(), + media_type="text/event-stream", + background=_global_state.tokenizer_manager.create_abort_task(obj), + ) + else: + try: + ret = await _global_state.tokenizer_manager.generate_request( + obj, request + ).__anext__() + return ret + except ValueError as e: + logger.error(f"[http_server] Error: {e}") + return _create_error_response(e) + + +@app.api_route("/generate_from_file", methods=["POST"]) +async def generate_from_file_request(file: UploadFile, request: Request): + """Handle a generate request, this is purely to work with input_embeds.""" + content = await file.read() + input_embeds = json.loads(content.decode("utf-8")) + + obj = GenerateReqInput( + input_embeds=input_embeds, + sampling_params={ + "temperature": 0.0, + "max_new_tokens": 512, + }, + ) + + try: + ret = await _global_state.tokenizer_manager.generate_request( + obj, request + ).__anext__() + return ret + except ValueError as e: + logger.error(f"Error: {e}") + return _create_error_response(e) + + +@app.api_route("/encode", methods=["POST", "PUT"]) +async def encode_request(obj: EmbeddingReqInput, request: Request): + """Handle an embedding request.""" + try: + ret = await _global_state.tokenizer_manager.generate_request( + obj, request + ).__anext__() + return ret + except ValueError as e: + return _create_error_response(e) + + +@app.api_route("/classify", methods=["POST", "PUT"]) +async def classify_request(obj: EmbeddingReqInput, request: Request): + """Handle a reward model request. Now the arguments and return values are the same as embedding models.""" + try: + ret = await _global_state.tokenizer_manager.generate_request( + obj, request + ).__anext__() + return ret + except ValueError as e: + return _create_error_response(e) + + +@app.api_route("/flush_cache", methods=["GET", "POST"]) +async def flush_cache(): + """Flush the radix cache.""" + ret = await _global_state.tokenizer_manager.flush_cache() + return Response( + content="Cache flushed.\nPlease check backend logs for more details. " + "(When there are running or waiting requests, the operation will not be performed.)\n", + status_code=200 if ret.success else HTTPStatus.BAD_REQUEST, + ) + + +@app.api_route("/clear_hicache_storage_backend", methods=["GET", "POST"]) +async def clear_hicache_storage_backend(): + """Clear the hierarchical cache storage backend.""" + ret = await _global_state.tokenizer_manager.clear_hicache_storage() + return Response( + content="Hierarchical cache storage backend cleared.\n", + status_code=200 if ret.success else HTTPStatus.BAD_REQUEST, + ) + + +@app.api_route("/start_profile", methods=["GET", "POST"]) +async def start_profile_async(obj: Optional[ProfileReqInput] = None): + """Start profiling.""" + if obj is None: + obj = ProfileReqInput() + + await _global_state.tokenizer_manager.start_profile( + output_dir=obj.output_dir, + start_step=obj.start_step, + num_steps=obj.num_steps, + activities=obj.activities, + with_stack=obj.with_stack, + record_shapes=obj.record_shapes, + profile_by_stage=obj.profile_by_stage, + ) + return Response( + content="Start profiling.\n", + status_code=200, + ) + + +@app.api_route("/stop_profile", methods=["GET", "POST"]) +async def stop_profile_async(): + """Stop profiling.""" + await _global_state.tokenizer_manager.stop_profile() + return Response( + content="Stop profiling. This will take some time.\n", + status_code=200, + ) + + +@app.api_route("/freeze_gc", methods=["GET", "POST"]) +async def freeze_gc_async(): + """ + See engine.freeze_gc for more details. + """ + await _global_state.tokenizer_manager.freeze_gc() + return Response( + content="Garbage collection frozen.\n", + status_code=200, + ) + + +@app.api_route("/start_expert_distribution_record", methods=["GET", "POST"]) +async def start_expert_distribution_record_async(): + """Start recording the expert distribution. Clear the previous record if any.""" + await _global_state.tokenizer_manager.start_expert_distribution_record() + return Response( + content="Start recording the expert distribution.\n", + status_code=200, + ) + + +@app.api_route("/stop_expert_distribution_record", methods=["GET", "POST"]) +async def stop_expert_distribution_record_async(): + """Stop recording the expert distribution.""" + await _global_state.tokenizer_manager.stop_expert_distribution_record() + return Response( + content="Stop recording the expert distribution.\n", + status_code=200, + ) + + +@app.api_route("/dump_expert_distribution_record", methods=["GET", "POST"]) +async def dump_expert_distribution_record_async(): + """Dump expert distribution record.""" + await _global_state.tokenizer_manager.dump_expert_distribution_record() + return Response( + content="Dump expert distribution record.\n", + status_code=200, + ) + + +@app.post("/update_weights_from_disk") +async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: Request): + """Update the weights from disk inplace without re-launching the server.""" + success, message, num_paused_requests = ( + await _global_state.tokenizer_manager.update_weights_from_disk(obj, request) + ) + + # Update weight version if provided and weights update was successful + if success and obj.weight_version is not None: + _update_weight_version_if_provided(obj.weight_version) + message += f" Weight version updated to {obj.weight_version}." + + content = { + "success": success, + "message": message, + "num_paused_requests": num_paused_requests, + } + if success: + return ORJSONResponse( + content, + status_code=HTTPStatus.OK, + ) + else: + return ORJSONResponse( + content, + status_code=HTTPStatus.BAD_REQUEST, + ) + + +@app.post("/init_weights_update_group") +async def init_weights_update_group( + obj: InitWeightsUpdateGroupReqInput, request: Request +): + """Initialize the parameter update group.""" + success, message = await _global_state.tokenizer_manager.init_weights_update_group( + obj, request + ) + content = {"success": success, "message": message} + if success: + return ORJSONResponse(content, status_code=200) + else: + return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) + + +@app.post("/update_weights_from_tensor") +async def update_weights_from_tensor( + obj: UpdateWeightsFromTensorReqInput, request: Request +): + """Update the weights from tensor inplace without re-launching the server. + Notes: + 1. Ensure that the model is on the correct device (e.g., GPU) before calling this endpoint. If the model is moved to the CPU unexpectedly, it may cause performance issues or runtime errors. + 2. HTTP will transmit only the metadata of the tensor, while the tensor itself will be directly copied to the model. + 3. Any binary data in the named tensors should be base64 encoded. + """ + + success, message = await _global_state.tokenizer_manager.update_weights_from_tensor( + obj, request + ) + + # Update weight version if provided and weights update was successful + if success and obj.weight_version is not None: + _update_weight_version_if_provided(obj.weight_version) + message += f" Weight version updated to {obj.weight_version}." + + content = {"success": success, "message": message} + return ORJSONResponse( + content, status_code=200 if success else HTTPStatus.BAD_REQUEST + ) + + +@app.post("/update_weights_from_distributed") +async def update_weights_from_distributed( + obj: UpdateWeightsFromDistributedReqInput, request: Request +): + """Update model parameter from distributed online.""" + success, message = ( + await _global_state.tokenizer_manager.update_weights_from_distributed( + obj, request + ) + ) + + # Update weight version if provided and weights update was successful + if success and obj.weight_version is not None: + _update_weight_version_if_provided(obj.weight_version) + message += f" Weight version updated to {obj.weight_version}." + + content = {"success": success, "message": message} + if success: + return ORJSONResponse(content, status_code=200) + else: + return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) + + +@app.post("/update_weight_version") +async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Request): + """Update the weight version. This operation requires no active requests.""" + if obj.abort_all_requests: + _global_state.tokenizer_manager.abort_request(abort_all=True) + + # Use a simple approach without the complex lock mechanism for now + # since weight_version update is a simple operation that doesn't affect model weights + try: + # Update the weight version in server args (the single source of truth) + _global_state.tokenizer_manager.server_args.weight_version = obj.new_version + + return ORJSONResponse( + { + "success": True, + "message": f"Weight version updated to {obj.new_version}", + "new_version": obj.new_version, + }, + status_code=HTTPStatus.OK, + ) + except Exception as e: + return ORJSONResponse( + { + "success": False, + "message": f"Failed to update weight version: {str(e)}", + }, + status_code=HTTPStatus.BAD_REQUEST, + ) + + +@app.api_route("/get_weights_by_name", methods=["GET", "POST"]) +async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request): + """Get model parameter by name.""" + try: + ret = await _global_state.tokenizer_manager.get_weights_by_name(obj, request) + if ret is None: + return _create_error_response("Get parameter by name failed") + else: + return ORJSONResponse(ret, status_code=200) + except Exception as e: + return _create_error_response(e) + + +@app.api_route("/release_memory_occupation", methods=["GET", "POST"]) +async def release_memory_occupation( + obj: ReleaseMemoryOccupationReqInput, request: Request +): + """Release GPU memory occupation temporarily.""" + try: + await _global_state.tokenizer_manager.release_memory_occupation(obj, request) + except Exception as e: + return _create_error_response(e) + + +@app.api_route("/resume_memory_occupation", methods=["GET", "POST"]) +async def resume_memory_occupation( + obj: ResumeMemoryOccupationReqInput, request: Request +): + """Resume GPU memory occupation.""" + try: + await _global_state.tokenizer_manager.resume_memory_occupation(obj, request) + except Exception as e: + return _create_error_response(e) + + +@app.api_route("/slow_down", methods=["GET", "POST"]) +async def slow_down(obj: SlowDownReqInput, request: Request): + """Slow down the system deliberately. Only for testing. Example scenario: + when we want to test performance of D in large-scale PD disaggregation and have no enough nodes for P, + we can use this to slow down D to let it have enough running sequences, and then disable slowdown + to let it run in full batch size. + """ + try: + await _global_state.tokenizer_manager.slow_down(obj, request) + except Exception as e: + return _create_error_response(e) + + +@app.api_route("/load_lora_adapter", methods=["POST"]) +async def load_lora_adapter(obj: LoadLoRAAdapterReqInput, request: Request): + """Load a new LoRA adapter without re-launching the server.""" + result = await _global_state.tokenizer_manager.load_lora_adapter(obj, request) + + if result.success: + return ORJSONResponse( + result, + status_code=HTTPStatus.OK, + ) + else: + return ORJSONResponse( + result, + status_code=HTTPStatus.BAD_REQUEST, + ) + + +@app.api_route("/unload_lora_adapter", methods=["POST"]) +async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInput, request: Request): + """Load a new LoRA adapter without re-launching the server.""" + result = await _global_state.tokenizer_manager.unload_lora_adapter(obj, request) + + if result.success: + return ORJSONResponse( + result, + status_code=HTTPStatus.OK, + ) + else: + return ORJSONResponse( + result, + status_code=HTTPStatus.BAD_REQUEST, + ) + + +@app.api_route("/open_session", methods=["GET", "POST"]) +async def open_session(obj: OpenSessionReqInput, request: Request): + """Open a session, and return its unique session id.""" + try: + session_id = await _global_state.tokenizer_manager.open_session(obj, request) + if session_id is None: + raise Exception( + "Failed to open the session. Check if a session with the same id is still open." + ) + return session_id + except Exception as e: + return _create_error_response(e) + + +@app.api_route("/close_session", methods=["GET", "POST"]) +async def close_session(obj: CloseSessionReqInput, request: Request): + """Close the session.""" + try: + await _global_state.tokenizer_manager.close_session(obj, request) + return Response(status_code=200) + except Exception as e: + return _create_error_response(e) + + +@app.api_route("/configure_logging", methods=["GET", "POST"]) +async def configure_logging(obj: ConfigureLoggingReq, request: Request): + """Configure the request logging options.""" + _global_state.tokenizer_manager.configure_logging(obj) + return Response(status_code=200) + + +@app.post("/abort_request") +async def abort_request(obj: AbortReq, request: Request): + """Abort a request.""" + try: + _global_state.tokenizer_manager.abort_request( + rid=obj.rid, abort_all=obj.abort_all + ) + return Response(status_code=200) + except Exception as e: + return _create_error_response(e) + + +@app.post("/parse_function_call") +async def parse_function_call_request(obj: ParseFunctionCallReq, request: Request): + """ + A native API endpoint to parse function calls from a text. + """ + # 1) Initialize the parser based on the request body + parser = FunctionCallParser(tools=obj.tools, tool_call_parser=obj.tool_call_parser) + + # 2) Call the non-stream parsing method (non-stream) + normal_text, calls = parser.parse_non_stream(obj.text) + + # 3) Organize the response content + response_data = { + "normal_text": normal_text, + "calls": [ + call.model_dump() for call in calls + ], # Convert pydantic objects to dictionaries + } + + return ORJSONResponse(content=response_data, status_code=200) + + +@app.post("/separate_reasoning") +async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Request): + """ + A native API endpoint to separate reasoning from a text. + """ + # 1) Initialize the parser based on the request body + parser = ReasoningParser(model_type=obj.reasoning_parser) + + # 2) Call the non-stream parsing method (non-stream) + reasoning_text, normal_text = parser.parse_non_stream(obj.text) + + # 3) Organize the response content + response_data = { + "reasoning_text": reasoning_text, + "text": normal_text, + } + + return ORJSONResponse(content=response_data, status_code=200) + + +@app.post("/pause_generation") +async def pause_generation(request: Request): + """Pause generation.""" + await _global_state.tokenizer_manager.pause_generation() + return ORJSONResponse( + content={"message": "Generation paused successfully.", "status": "ok"}, + status_code=200, + ) + + +@app.post("/continue_generation") +async def continue_generation(request: Request): + """Continue generation.""" + await _global_state.tokenizer_manager.continue_generation() + return ORJSONResponse( + content={"message": "Generation continued successfully.", "status": "ok"}, + status_code=200, + ) + + +##### OpenAI-compatible API endpoints ##### + + +@app.post("/v1/completions", dependencies=[Depends(validate_json_request)]) +async def openai_v1_completions(request: CompletionRequest, raw_request: Request): + """OpenAI-compatible text completion endpoint.""" + return await raw_request.app.state.openai_serving_completion.handle_request( + request, raw_request + ) + + +@app.post("/v1/chat/completions", dependencies=[Depends(validate_json_request)]) +async def openai_v1_chat_completions( + request: ChatCompletionRequest, raw_request: Request +): + """OpenAI-compatible chat completion endpoint.""" + return await raw_request.app.state.openai_serving_chat.handle_request( + request, raw_request + ) + + +@app.post( + "/v1/embeddings", + response_class=ORJSONResponse, + dependencies=[Depends(validate_json_request)], +) +async def openai_v1_embeddings(request: EmbeddingRequest, raw_request: Request): + """OpenAI-compatible embeddings endpoint.""" + return await raw_request.app.state.openai_serving_embedding.handle_request( + request, raw_request + ) + + +@app.get("/v1/models", response_class=ORJSONResponse) +async def available_models(): + """Show available models. OpenAI-compatible endpoint.""" + served_model_names = [_global_state.tokenizer_manager.served_model_name] + model_cards = [] + for served_model_name in served_model_names: + model_cards.append( + ModelCard( + id=served_model_name, + root=served_model_name, + max_model_len=_global_state.tokenizer_manager.model_config.context_len, + ) + ) + return ModelList(data=model_cards) + + +@app.get("/v1/models/{model:path}", response_class=ORJSONResponse) +async def retrieve_model(model: str): + """Retrieves a model instance, providing basic information about the model.""" + served_model_names = [_global_state.tokenizer_manager.served_model_name] + + if model not in served_model_names: + return ORJSONResponse( + status_code=404, + content={ + "error": { + "message": f"The model '{model}' does not exist", + "type": "invalid_request_error", + "param": "model", + "code": "model_not_found", + } + }, + ) + + return ModelCard( + id=model, + root=model, + max_model_len=_global_state.tokenizer_manager.model_config.context_len, + ) + + +@app.post("/v1/score", dependencies=[Depends(validate_json_request)]) +async def v1_score_request(request: ScoringRequest, raw_request: Request): + """Endpoint for the decoder-only scoring API. See Engine.score() for detailed documentation.""" + return await raw_request.app.state.openai_serving_score.handle_request( + request, raw_request + ) + + +@app.post("/v1/responses", dependencies=[Depends(validate_json_request)]) +async def v1_responses_request(request: dict, raw_request: Request): + """Endpoint for the responses API with reasoning support.""" + + request_obj = ResponsesRequest(**request) + result = await raw_request.app.state.openai_serving_responses.create_responses( + request_obj, raw_request + ) + + # Handle streaming responses + if isinstance(result, AsyncGenerator): + return StreamingResponse( + result, + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, + ) + + return result + + +@app.get("/v1/responses/{response_id}") +async def v1_retrieve_responses(response_id: str, raw_request: Request): + """Retrieve a response by ID.""" + return await raw_request.app.state.openai_serving_responses.retrieve_responses( + response_id + ) + + +@app.post("/v1/responses/{response_id}/cancel") +async def v1_cancel_responses(response_id: str, raw_request: Request): + """Cancel a background response.""" + return await raw_request.app.state.openai_serving_responses.cancel_responses( + response_id + ) + + +@app.api_route( + "/v1/rerank", methods=["POST", "PUT"], dependencies=[Depends(validate_json_request)] +) +async def v1_rerank_request(request: V1RerankReqInput, raw_request: Request): + """Endpoint for reranking documents based on query relevance.""" + return await raw_request.app.state.openai_serving_rerank.handle_request( + request, raw_request + ) + + +## SageMaker API +@app.get("/ping") +async def sagemaker_health() -> Response: + """Check the health of the http server.""" + return Response(status_code=200) + + +@app.post("/invocations") +async def sagemaker_chat_completions( + request: ChatCompletionRequest, raw_request: Request +): + """OpenAI-compatible chat completion endpoint.""" + return await raw_request.app.state.openai_serving_chat.handle_request( + request, raw_request + ) + + +## Vertex AI API +@app.post(os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate")) +async def vertex_generate(vertex_req: VertexGenerateReqInput, raw_request: Request): + if not vertex_req.instances: + return [] + inputs = {} + for input_key in ("text", "input_ids", "input_embeds"): + if vertex_req.instances[0].get(input_key): + inputs[input_key] = [ + instance.get(input_key) for instance in vertex_req.instances + ] + break + image_data = [ + instance.get("image_data") + for instance in vertex_req.instances + if instance.get("image_data") is not None + ] or None + req = GenerateReqInput( + **inputs, + image_data=image_data, + **(vertex_req.parameters or {}), + ) + ret = await generate_request(req, raw_request) + if isinstance(ret, Response): + return ret + return ORJSONResponse({"predictions": ret}) + + +def _update_weight_version_if_provided(weight_version: Optional[str]) -> None: + """Update weight version if provided.""" + if weight_version is not None: + _global_state.tokenizer_manager.server_args.weight_version = weight_version + + +def _create_error_response(e): + return ORJSONResponse( + {"error": {"message": str(e)}}, status_code=HTTPStatus.BAD_REQUEST + ) + + +def launch_server( + server_args: ServerArgs, + pipe_finish_writer: Optional[multiprocessing.connection.Connection] = None, + launch_callback: Optional[Callable[[], None]] = None, +): + """ + Launch SRT (SGLang Runtime) Server. + + The SRT server consists of an HTTP server and an SRT engine. + + - HTTP server: A FastAPI server that routes requests to the engine. + - The engine consists of three components: + 1. TokenizerManager: Tokenizes the requests and sends them to the scheduler. + 2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager. + 3. DetokenizerManager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager. + + Note: + 1. The HTTP server, Engine, and TokenizerManager both run in the main process. + 2. Inter-process communication is done through IPC (each process uses a different port) via the ZMQ library. + """ + if server_args.tokenizer_worker_num > 1: + setproctitle.setproctitle(f"sglang::http_server/multi_tokenizer_router") + port_args = PortArgs.init_new(server_args) + port_args.tokenizer_worker_ipc_name = ( + f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}" + ) + tokenizer_manager, template_manager, scheduler_info = _launch_subprocesses( + server_args=server_args, port_args=port_args + ) + else: + setproctitle.setproctitle(f"sglang::http_server/tokenizer_manager") + tokenizer_manager, template_manager, scheduler_info = _launch_subprocesses( + server_args=server_args, + ) + + set_global_state( + _GlobalState( + tokenizer_manager=tokenizer_manager, + template_manager=template_manager, + scheduler_info=scheduler_info, + ) + ) + + if server_args.tokenizer_worker_num > 1: + multi_tokenizer_args_shm = write_data_for_multi_tokenizer( + port_args, + server_args, + scheduler_info, + ) + else: + # Add api key authorization + if server_args.api_key: + add_api_key_middleware(app, server_args.api_key) + + # Add prometheus middleware + if server_args.enable_metrics: + add_prometheus_middleware(app) + enable_func_timer() + + # Send a warmup request - we will create the thread launch it + # in the lifespan after all other warmups have fired. + warmup_thread = threading.Thread( + target=_wait_and_warmup, + args=( + server_args, + pipe_finish_writer, + launch_callback, + ), + ) + app.warmup_thread = warmup_thread + + try: + # Update logging configs + set_uvicorn_logging_configs() + app.server_args = server_args + # Listen for HTTP requests + if server_args.tokenizer_worker_num > 1: + from uvicorn.config import LOGGING_CONFIG + + LOGGING_CONFIG["loggers"]["sglang.srt.entrypoints.http_server"] = { + "handlers": ["default"], + "level": "INFO", + "propagate": False, + } + + monkey_patch_uvicorn_multiprocessing() + + uvicorn.run( + "sglang.srt.entrypoints.http_server:app", + host=server_args.host, + port=server_args.port, + log_level=server_args.log_level_http or server_args.log_level, + timeout_keep_alive=5, + loop="uvloop", + workers=server_args.tokenizer_worker_num, + ) + else: + app.is_single_tokenizer_mode = True + uvicorn.run( + app, + host=server_args.host, + port=server_args.port, + log_level=server_args.log_level_http or server_args.log_level, + timeout_keep_alive=5, + loop="uvloop", + ) + finally: + if server_args.tokenizer_worker_num > 1: + multi_tokenizer_args_shm.unlink() + _global_state.tokenizer_manager.socket_mapping.clear_all_sockets() + else: + warmup_thread.join() + + +def _execute_server_warmup( + server_args: ServerArgs, + pipe_finish_writer: Optional[multiprocessing.connection.Connection], +): + headers = {} + url = server_args.url() + if server_args.api_key: + headers["Authorization"] = f"Bearer {server_args.api_key}" + + # Wait until the server is launched + success = False + for _ in range(120): + time.sleep(1) + try: + res = requests.get(url + "/get_model_info", timeout=5, headers=headers) + assert res.status_code == 200, f"{res=}, {res.text=}" + success = True + break + except (AssertionError, requests.exceptions.RequestException): + last_traceback = get_exception_traceback() + pass + + if not success: + if pipe_finish_writer is not None: + pipe_finish_writer.send(last_traceback) + logger.error(f"Initialization failed. warmup error: {last_traceback}") + kill_process_tree(os.getpid()) + return success + + model_info = res.json() + + # Send a warmup request + request_name = "/generate" if model_info["is_generation"] else "/encode" + max_new_tokens = 8 if model_info["is_generation"] else 1 + json_data = { + "sampling_params": { + "temperature": 0, + "max_new_tokens": max_new_tokens, + }, + } + if server_args.skip_tokenizer_init: + json_data["input_ids"] = [[10, 11, 12] for _ in range(server_args.dp_size)] + # TODO Workaround the bug that embedding errors for list of size 1 + if server_args.dp_size == 1: + json_data["input_ids"] = json_data["input_ids"][0] + else: + json_data["text"] = ["The capital city of France is"] * server_args.dp_size + # TODO Workaround the bug that embedding errors for list of size 1 + if server_args.dp_size == 1: + json_data["text"] = json_data["text"][0] + + # Debug dumping + if server_args.debug_tensor_dump_input_file: + json_data.pop("text", None) + json_data["input_ids"] = np.load( + server_args.debug_tensor_dump_input_file + ).tolist() + json_data["sampling_params"]["max_new_tokens"] = 0 + + try: + if server_args.disaggregation_mode == "null": + res = requests.post( + url + request_name, + json=json_data, + headers=headers, + timeout=600, + ) + assert res.status_code == 200, f"{res}" + _global_state.tokenizer_manager.server_status = ServerStatus.Up + + else: + logger.info(f"Start of pd disaggregation warmup ...") + json_data = { + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": 8, + "ignore_eos": True, + }, + "bootstrap_host": [FAKE_BOOTSTRAP_HOST] * server_args.dp_size, + # This is a hack to ensure fake transfer is enabled during prefill warmup + # ensure each dp rank has a unique bootstrap_room during prefill warmup + "bootstrap_room": [ + i * (2**63 // server_args.dp_size) + (i % server_args.tp_size) + for i in range(server_args.dp_size) + ], + "input_ids": [[0, 1, 2, 3]] * server_args.dp_size, + } + res = requests.post( + url + request_name, + json=json_data, + headers=headers, + timeout=1800, # because of deep gemm precache is very long if not precache. + ) + if res.status_code == 200: + logger.info( + f"End of prefill disaggregation mode warmup with status {res.status_code}, resp: {res.json()}" + ) + _global_state.tokenizer_manager.server_status = ServerStatus.Up + else: + logger.info( + "Prefill disaggregation mode warm Up Failed, status code: {}".format( + res.status_code + ) + ) + _global_state.tokenizer_manager.server_status = ServerStatus.UnHealthy + + except Exception: + last_traceback = get_exception_traceback() + if pipe_finish_writer is not None: + pipe_finish_writer.send(last_traceback) + logger.error(f"Initialization failed. warmup error: {last_traceback}") + kill_process_tree(os.getpid()) + return False + + # Debug print + # logger.info(f"warmup request returns: {res.json()=}") + return success + + +def _wait_and_warmup( + server_args: ServerArgs, + pipe_finish_writer: Optional[multiprocessing.connection.Connection], + launch_callback: Optional[Callable[[], None]] = None, +): + if not server_args.skip_server_warmup: + if not _execute_server_warmup( + server_args, + pipe_finish_writer, + ): + return + else: + _global_state.tokenizer_manager.server_status = ServerStatus.Up + + logger.info("The server is fired up and ready to roll!") + + if pipe_finish_writer is not None: + pipe_finish_writer.send("ready") + + if server_args.delete_ckpt_after_loading: + delete_directory(server_args.model_path) + + if server_args.debug_tensor_dump_input_file: + kill_process_tree(os.getpid()) + + if launch_callback is not None: + launch_callback() diff --git a/python/sglang/srt/entrypoints/http_server_engine.py b/python/sglang/srt/entrypoints/http_server_engine.py new file mode 100644 index 000000000..d1db80d65 --- /dev/null +++ b/python/sglang/srt/entrypoints/http_server_engine.py @@ -0,0 +1,142 @@ +import copy +import dataclasses +import multiprocessing +import pickle +import threading +import time +from typing import Any, Dict, List, Optional, Tuple, Union + +import pybase64 +import requests +import torch +import torch.distributed as dist + +from sglang.srt.entrypoints.EngineBase import EngineBase +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import MultiprocessingSerializer, kill_process_tree + + +def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process: + + p = multiprocessing.Process(target=launch_server, args=(server_args,)) + p.start() + + base_url = server_args.url() + timeout = 300.0 # Increased timeout to 5 minutes for downloading large models + start_time = time.perf_counter() + + with requests.Session() as session: + while time.perf_counter() - start_time < timeout: + try: + headers = { + "Content-Type": "application/json; charset=utf-8", + "Authorization": f"Bearer {server_args.api_key}", + } + response = session.get(f"{base_url}/health_generate", headers=headers) + if response.status_code == 200: + return p + except requests.RequestException: + pass + + if not p.is_alive(): + raise Exception("Server process terminated unexpectedly.") + + time.sleep(2) + + p.terminate() + raise TimeoutError("Server failed to start within the timeout period.") + + +class HttpServerEngineAdapter(EngineBase): + """ + You can use this class to launch a server from a VerlEngine instance. + We recommend using this class only you need to use http server. + Otherwise, you can use Engine directly. + """ + + def __init__(self, **kwargs): + self.server_args = ServerArgs(**kwargs) + print( + f"Launch HttpServerEngineAdapter at: {self.server_args.host}:{self.server_args.port}" + ) + self.process = launch_server_process(self.server_args) + + def _make_request(self, endpoint: str, payload: Optional[dict] = None): + """Make a POST request to the specified endpoint with the given payload. + Args: + endpoint: The API endpoint to call + payload: The JSON payload to send (default: empty dict) + Returns: + The JSON response from the server + """ + url = f"http://{self.server_args.host}:{self.server_args.port}/{endpoint}" + response = requests.post(url, json=payload or {}) + response.raise_for_status() + return response.json() + + def update_weights_from_tensor( + self, + named_tensors: List[Tuple[str, torch.Tensor]], + load_format: Optional[str] = None, + flush_cache: bool = False, + ): + """ + Update model weights from tensor data. The HTTP server will only post meta data, and the real weights will be copied directly from GPUs. + Note: The model should be on GPUs rather than CPU for this functionality to work properly. + If you encounter issues, ensure your model is loaded on GPU devices rather than CPU. + """ + + return self._make_request( + "update_weights_from_tensor", + { + "serialized_named_tensors": [ + MultiprocessingSerializer.serialize(named_tensors, output_str=True) + for _ in range(self.server_args.tp_size) + ], + "load_format": load_format, + "flush_cache": flush_cache, + }, + ) + + def shutdown(self): + kill_process_tree(self.process.pid) + + def generate( + self, + prompt=None, + sampling_params=None, + input_ids=None, + image_data=None, + return_logprob=False, + logprob_start_len=None, + top_logprobs_num=None, + token_ids_logprob=None, + lora_path=None, + custom_logit_processor=None, + ): + payload = { + "text": prompt, + "sampling_params": sampling_params, + "input_ids": input_ids, + "image_data": image_data, + "return_logprob": return_logprob, + "logprob_start_len": logprob_start_len, + "top_logprobs_num": top_logprobs_num, + "token_ids_logprob": token_ids_logprob, + "lora_path": lora_path, + "custom_logit_processor": custom_logit_processor, + } + # Filter out None values + payload = {k: v for k, v in payload.items() if v is not None} + + return self._make_request("generate", payload) + + def release_memory_occupation(self): + return self._make_request("release_memory_occupation") + + def resume_memory_occupation(self): + return self._make_request("resume_memory_occupation") + + def flush_cache(self): + return self._make_request("flush_cache") diff --git a/python/sglang/srt/entrypoints/openai/__init__.py b/python/sglang/srt/entrypoints/openai/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py new file mode 100644 index 000000000..f73e67d0b --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -0,0 +1,929 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Pydantic models for OpenAI API protocol""" + +import time +import uuid +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, TypeAlias, Union + +from openai.types.responses import ( + ResponseFunctionToolCall, + ResponseInputItemParam, + ResponseOutputItem, + ResponseReasoningItem, +) +from openai.types.responses.response import ToolChoice +from openai.types.responses.tool import Tool +from pydantic import ( + BaseModel, + Field, + field_validator, + model_serializer, + model_validator, +) +from typing_extensions import Literal + +DEFAULT_MODEL_NAME = "default" + + +class ModelCard(BaseModel): + """Model cards.""" + + id: str + object: str = "model" + created: int = Field(default_factory=lambda: int(time.time())) + owned_by: str = "sglang" + root: Optional[str] = None + max_model_len: Optional[int] = None + + +class ModelList(BaseModel): + """Model list consists of model cards.""" + + object: str = "list" + data: List[ModelCard] = Field(default_factory=list) + + +class ErrorResponse(BaseModel): + object: str = "error" + message: str + type: str + param: Optional[str] = None + code: int + + +class LogProbs(BaseModel): + text_offset: List[int] = Field(default_factory=list) + token_logprobs: List[Optional[float]] = Field(default_factory=list) + tokens: List[str] = Field(default_factory=list) + top_logprobs: List[Optional[Dict[str, float]]] = Field(default_factory=list) + + +class TopLogprob(BaseModel): + token: str + bytes: List[int] + logprob: float + + +class ChatCompletionTokenLogprob(BaseModel): + token: str + bytes: List[int] + logprob: float + top_logprobs: List[TopLogprob] + + +class ChoiceLogprobs(BaseModel): + # build for v1/chat/completions response + content: List[ChatCompletionTokenLogprob] + + +class UsageInfo(BaseModel): + prompt_tokens: int = 0 + total_tokens: int = 0 + completion_tokens: Optional[int] = 0 + # only used to return cached tokens when --enable-cache-report is set + prompt_tokens_details: Optional[Dict[str, int]] = None + reasoning_tokens: Optional[int] = 0 + + +class StreamOptions(BaseModel): + include_usage: Optional[bool] = False + + +class JsonSchemaResponseFormat(BaseModel): + name: str + description: Optional[str] = None + # use alias to workaround pydantic conflict + schema_: Optional[Dict[str, object]] = Field(alias="schema", default=None) + strict: Optional[bool] = False + + +class ResponseFormat(BaseModel): + type: Literal["text", "json_object", "json_schema"] + json_schema: Optional[JsonSchemaResponseFormat] = None + + +class StructuresResponseFormat(BaseModel): + begin: str + schema_: Optional[Dict[str, object]] = Field(alias="schema", default=None) + end: str + + +class StructuralTagResponseFormat(BaseModel): + type: Literal["structural_tag"] + structures: List[StructuresResponseFormat] + triggers: List[str] + + +class FileRequest(BaseModel): + # https://platform.openai.com/docs/api-reference/files/create + file: bytes # The File object (not file name) to be uploaded + purpose: str = ( + "batch" # The intended purpose of the uploaded file, default is "batch" + ) + + +class FileResponse(BaseModel): + id: str + object: str = "file" + bytes: int + created_at: int + filename: str + purpose: str + + +class FileDeleteResponse(BaseModel): + id: str + object: str = "file" + deleted: bool + + +class BatchRequest(BaseModel): + input_file_id: ( + str # The ID of an uploaded file that contains requests for the new batch + ) + endpoint: str # The endpoint to be used for all requests in the batch + completion_window: str # The time frame within which the batch should be processed + metadata: Optional[dict] = None # Optional custom metadata for the batch + + +class BatchResponse(BaseModel): + id: str + object: str = "batch" + endpoint: str + errors: Optional[dict] = None + input_file_id: str + completion_window: str + status: str = "validating" + output_file_id: Optional[str] = None + error_file_id: Optional[str] = None + created_at: int + in_progress_at: Optional[int] = None + expires_at: Optional[int] = None + finalizing_at: Optional[int] = None + completed_at: Optional[int] = None + failed_at: Optional[int] = None + expired_at: Optional[int] = None + cancelling_at: Optional[int] = None + cancelled_at: Optional[int] = None + request_counts: Optional[dict] = None + metadata: Optional[dict] = None + + +class CompletionRequest(BaseModel): + # Ordered by official OpenAI API documentation + # https://platform.openai.com/docs/api-reference/completions/create + model: str = DEFAULT_MODEL_NAME + prompt: Union[List[int], List[List[int]], str, List[str]] + best_of: Optional[int] = None + echo: bool = False + frequency_penalty: float = 0.0 + logit_bias: Optional[Dict[str, float]] = None + logprobs: Optional[int] = None + max_tokens: int = 16 + n: int = 1 + presence_penalty: float = 0.0 + seed: Optional[int] = None + stop: Optional[Union[str, List[str]]] = None + stream: bool = False + stream_options: Optional[StreamOptions] = None + suffix: Optional[str] = None + temperature: float = 1.0 + top_p: float = 1.0 + user: Optional[str] = None + return_hidden_states: bool = False + + # Extra parameters for SRT backend only and will be ignored by OpenAI models. + top_k: int = -1 + min_p: float = 0.0 + min_tokens: int = 0 + json_schema: Optional[str] = None + regex: Optional[str] = None + ebnf: Optional[str] = None + repetition_penalty: float = 1.0 + stop_token_ids: Optional[List[int]] = None + no_stop_trim: bool = False + ignore_eos: bool = False + skip_special_tokens: bool = True + lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None + session_params: Optional[Dict] = None + response_format: Optional[Union[ResponseFormat, StructuralTagResponseFormat]] = None + + # For PD disaggregation + bootstrap_host: Optional[Union[List[str], str]] = None + bootstrap_port: Optional[Union[List[Optional[int]], int]] = None + bootstrap_room: Optional[Union[List[int], int]] = None + + # For request id + rid: Optional[Union[List[str], str]] = None + + @field_validator("max_tokens") + @classmethod + def validate_max_tokens_positive(cls, v): + if v is not None and v <= 0: + raise ValueError("max_tokens must be positive") + return v + + +class CompletionResponseChoice(BaseModel): + index: int + text: str + logprobs: Optional[LogProbs] = None + finish_reason: Optional[Literal["stop", "length", "content_filter", "abort"]] = None + matched_stop: Union[None, int, str] = None + hidden_states: Optional[object] = None + + @model_serializer(mode="wrap") + def _serialize(self, handler): + data = handler(self) + if self.hidden_states is None: + data.pop("hidden_states", None) + return data + + +class CompletionResponse(BaseModel): + id: str + object: str = "text_completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[CompletionResponseChoice] + usage: UsageInfo + metadata: Optional[Dict[str, Any]] = None + + +class CompletionResponseStreamChoice(BaseModel): + index: int + text: str + logprobs: Optional[LogProbs] = None + finish_reason: Optional[Literal["stop", "length", "content_filter", "abort"]] = None + matched_stop: Union[None, int, str] = None + hidden_states: Optional[object] = None + + @model_serializer(mode="wrap") + def _serialize(self, handler): + data = handler(self) + if self.hidden_states is None: + data.pop("hidden_states", None) + return data + + +class CompletionStreamResponse(BaseModel): + id: str + object: str = "text_completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[CompletionResponseStreamChoice] + usage: Optional[UsageInfo] = None + + +class ChatCompletionMessageContentTextPart(BaseModel): + type: Literal["text"] + text: str + + +class ChatCompletionMessageContentImageURL(BaseModel): + url: str + detail: Optional[Literal["auto", "low", "high"]] = "auto" + + +class ChatCompletionMessageContentVideoURL(BaseModel): + url: str + + +class ChatCompletionMessageContentAudioURL(BaseModel): + url: str + + +class ChatCompletionMessageContentImagePart(BaseModel): + type: Literal["image_url"] + image_url: ChatCompletionMessageContentImageURL + modalities: Optional[Literal["image", "multi-images", "video"]] = "image" + + +class ChatCompletionMessageContentVideoPart(BaseModel): + type: Literal["video_url"] + video_url: ChatCompletionMessageContentVideoURL + + +class ChatCompletionMessageContentAudioPart(BaseModel): + type: Literal["audio_url"] + audio_url: ChatCompletionMessageContentAudioURL + + +ChatCompletionMessageContentPart = Union[ + ChatCompletionMessageContentTextPart, + ChatCompletionMessageContentImagePart, + ChatCompletionMessageContentVideoPart, + ChatCompletionMessageContentAudioPart, +] + + +class FunctionResponse(BaseModel): + """Function response.""" + + name: Optional[str] = None + arguments: Optional[str] = None + + +class ToolCall(BaseModel): + """Tool call response.""" + + id: Optional[str] = None + index: Optional[int] = None + type: Literal["function"] = "function" + function: FunctionResponse + + +class ChatCompletionMessageGenericParam(BaseModel): + role: Literal["system", "assistant", "tool", "function"] + content: Union[str, List[ChatCompletionMessageContentTextPart], None] = Field( + default=None + ) + tool_call_id: Optional[str] = None + name: Optional[str] = None + reasoning_content: Optional[str] = None + tool_calls: Optional[List[ToolCall]] = Field(default=None, examples=[None]) + + @field_validator("role", mode="before") + @classmethod + def _normalize_role(cls, v): + if isinstance(v, str): + v_lower = v.lower() + if v_lower not in {"system", "assistant", "tool", "function"}: + raise ValueError( + "'role' must be one of 'system', 'assistant', 'tool', or 'function' (case-insensitive)." + ) + return v_lower + raise ValueError("'role' must be a string") + + +class ChatCompletionMessageUserParam(BaseModel): + role: Literal["user"] + content: Union[str, List[ChatCompletionMessageContentPart]] + + +ChatCompletionMessageParam = Union[ + ChatCompletionMessageGenericParam, ChatCompletionMessageUserParam +] + + +class Function(BaseModel): + """Function descriptions.""" + + description: Optional[str] = Field(default=None, examples=[None]) + name: Optional[str] = None + parameters: Optional[object] = None + strict: bool = False + + +class Tool(BaseModel): + """Function wrapper.""" + + type: str = Field(default="function", examples=["function"]) + function: Function + + +class ToolChoiceFuncName(BaseModel): + """The name of tool choice function.""" + + name: Optional[str] = None + + +class ToolChoice(BaseModel): + """The tool choice definition.""" + + function: ToolChoiceFuncName + type: Literal["function"] = Field(default="function", examples=["function"]) + + +class ChatCompletionRequest(BaseModel): + # Ordered by official OpenAI API documentation + # https://platform.openai.com/docs/api-reference/chat/create + messages: List[ChatCompletionMessageParam] + model: str = DEFAULT_MODEL_NAME + frequency_penalty: float = 0.0 + logit_bias: Optional[Dict[str, float]] = None + logprobs: bool = False + top_logprobs: Optional[int] = None + max_tokens: Optional[int] = Field( + default=None, + deprecated="max_tokens is deprecated in favor of the max_completion_tokens field", + description="The maximum number of tokens that can be generated in the chat completion. ", + ) + max_completion_tokens: Optional[int] = Field( + default=None, + description="The maximum number of completion tokens for a chat completion request, " + "including visible output tokens and reasoning tokens. Input tokens are not included. ", + ) + n: int = 1 + presence_penalty: float = 0.0 + response_format: Optional[Union[ResponseFormat, StructuralTagResponseFormat]] = None + seed: Optional[int] = None + stop: Optional[Union[str, List[str]]] = None + stream: bool = False + stream_options: Optional[StreamOptions] = None + temperature: float = 0.7 + top_p: float = 1.0 + user: Optional[str] = None + tools: Optional[List[Tool]] = Field(default=None, examples=[None]) + tool_choice: Union[ToolChoice, Literal["auto", "required", "none"]] = Field( + default="auto", examples=["none"] + ) # noqa + return_hidden_states: bool = False + reasoning_effort: Optional[Literal["low", "medium", "high"]] = Field( + default="medium", + description="Constrains effort on reasoning for reasoning models. " + "'low' is the least effort, 'high' is the most effort. Reducing reasoning effort can " + "result in faster responses and fewer tokens used on reasoning in a response. " + "Currently only supported for OpenAI models.", + ) + + @model_validator(mode="before") + @classmethod + def set_tool_choice_default(cls, values): + if values.get("tool_choice") is None: + if values.get("tools") is None: + values["tool_choice"] = "none" + else: + values["tool_choice"] = "auto" + return values + + @model_validator(mode="before") + @classmethod + def normalize_reasoning_inputs(cls, values: Dict): + r = values.get("reasoning") + if r is None: + return values + + if isinstance(r, dict): + effort = r.get("effort") or r.get("reasoning_effort") + if effort in {"low", "medium", "high"}: + values["reasoning_effort"] = effort + + enabled = ( + r.get("enabled") + if r.get("enabled") is not None + else r.get("enable", False) + ) + if isinstance(enabled, str): + enabled = enabled.strip().lower() in {"1", "true", "yes", "y", "on"} + if enabled: + ctk = values.get("chat_template_kwargs") + if not isinstance(ctk, dict): + ctk = {} + ctk.setdefault("thinking", True) + values["chat_template_kwargs"] = ctk + + return values + + @model_validator(mode="before") + @classmethod + def set_json_schema(cls, values): + response_format = values.get("response_format") + if not response_format: + return values + + if response_format.get("type") != "json_schema": + return values + + schema = response_format.pop("schema", None) + json_schema = response_format.get("json_schema") + + if json_schema: + return values + + if schema: + name_ = schema.get("title", "Schema") + strict_ = False + if "properties" in schema and "strict" in schema["properties"]: + item = schema["properties"].pop("strict", None) + if item and item.get("default", False): + strict_ = True + + response_format["json_schema"] = { + "name": name_, + "schema": schema, + "strict": strict_, + } + + return values + + # Extra parameters for SRT backend only and will be ignored by OpenAI models. + top_k: int = -1 + min_p: float = 0.0 + min_tokens: int = 0 + regex: Optional[str] = None + ebnf: Optional[str] = None + repetition_penalty: float = 1.0 + stop_token_ids: Optional[List[int]] = None + no_stop_trim: bool = False + ignore_eos: bool = False + continue_final_message: bool = False + skip_special_tokens: bool = True + lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None + session_params: Optional[Dict] = None + separate_reasoning: bool = True + stream_reasoning: bool = True + chat_template_kwargs: Optional[Dict] = None + + # For request id + rid: Optional[Union[List[str], str]] = None + + # For PD disaggregation + bootstrap_host: Optional[Union[List[str], str]] = None + bootstrap_port: Optional[Union[List[Optional[int]], int]] = None + bootstrap_room: Optional[Union[List[int], int]] = None + + +class ChatMessage(BaseModel): + role: Optional[str] = None + content: Optional[str] = None + reasoning_content: Optional[str] = None + tool_calls: Optional[List[ToolCall]] = Field(default=None, examples=[None]) + + +class ChatCompletionResponseChoice(BaseModel): + index: int + message: ChatMessage + logprobs: Optional[Union[LogProbs, ChoiceLogprobs]] = None + finish_reason: Optional[ + Literal[ + "stop", "length", "tool_calls", "content_filter", "function_call", "abort" + ] + ] = None + matched_stop: Union[None, int, str] = None + hidden_states: Optional[object] = None + + @model_serializer(mode="wrap") + def _serialize(self, handler): + data = handler(self) + if self.hidden_states is None: + data.pop("hidden_states", None) + return data + + +class ChatCompletionResponse(BaseModel): + id: str + object: str = "chat.completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[ChatCompletionResponseChoice] + usage: UsageInfo + metadata: Optional[Dict[str, Any]] = None + + +class DeltaMessage(BaseModel): + role: Optional[str] = None + content: Optional[str] = None + reasoning_content: Optional[str] = None + tool_calls: Optional[List[ToolCall]] = Field(default=None, examples=[None]) + hidden_states: Optional[object] = None + + @model_serializer(mode="wrap") + def _serialize(self, handler): + data = handler(self) + if self.hidden_states is None: + data.pop("hidden_states", None) + return data + + +class ChatCompletionResponseStreamChoice(BaseModel): + index: int + delta: DeltaMessage + logprobs: Optional[Union[LogProbs, ChoiceLogprobs]] = None + finish_reason: Optional[ + Literal[ + "stop", "length", "tool_calls", "content_filter", "function_call", "abort" + ] + ] = None + matched_stop: Union[None, int, str] = None + + +class ChatCompletionStreamResponse(BaseModel): + id: str + object: str = "chat.completion.chunk" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[ChatCompletionResponseStreamChoice] + usage: Optional[UsageInfo] = None + + +class MultimodalEmbeddingInput(BaseModel): + text: Optional[str] = None + image: Optional[str] = None + + +EmbeddingInput = Union[ + List[int], List[List[int]], str, List[str], List[MultimodalEmbeddingInput] +] + + +class EmbeddingRequest(BaseModel): + # Ordered by official OpenAI API documentation + # https://platform.openai.com/docs/api-reference/embeddings/create + input: EmbeddingInput + model: str = DEFAULT_MODEL_NAME + encoding_format: str = "float" + dimensions: Optional[int] = None + user: Optional[str] = None + + # The request id. + rid: Optional[Union[List[str], str]] = None + + +class EmbeddingObject(BaseModel): + embedding: List[float] + index: int + object: str = "embedding" + + +class EmbeddingResponse(BaseModel): + data: List[EmbeddingObject] + model: str + object: str = "list" + usage: Optional[UsageInfo] = None + + +class ScoringRequest(BaseModel): + query: Optional[Union[str, List[int]]] = ( + None # Query text or pre-tokenized token IDs + ) + items: Optional[Union[str, List[str], List[List[int]]]] = ( + None # Item text(s) or pre-tokenized token IDs + ) + label_token_ids: Optional[List[int]] = ( + None # Token IDs to compute probabilities for + ) + apply_softmax: bool = False + item_first: bool = False + model: str = DEFAULT_MODEL_NAME + + +class ScoringResponse(BaseModel): + scores: List[ + List[float] + ] # List of lists of probabilities, each in the order of label_token_ids + model: str + usage: Optional[UsageInfo] = None + object: str = "scoring" + + +class V1RerankReqInput(BaseModel): + query: str + documents: List[str] + + +class RerankResponse(BaseModel): + score: float + document: str + index: int + meta_info: Optional[dict] = None + + +OpenAIServingRequest = Union[ + ChatCompletionRequest, + CompletionRequest, + EmbeddingRequest, + ScoringRequest, + V1RerankReqInput, +] + + +# Response API protocol definitions +class ResponseReasoningParam(BaseModel): + """Reasoning parameters for responses.""" + + effort: Optional[Literal["low", "medium", "high"]] = Field( + default="medium", + description="Constrains effort on reasoning for reasoning models.", + ) + + +class ResponseTool(BaseModel): + """Tool definition for responses.""" + + type: Literal["web_search_preview", "code_interpreter"] = Field( + description="Type of tool to enable" + ) + + +ResponseInputOutputItem: TypeAlias = Union[ + ResponseInputItemParam, + "ResponseReasoningItem", + ResponseFunctionToolCall, +] + + +class ResponsesRequest(BaseModel): + """Request body for v1/responses endpoint.""" + + # Core OpenAI API fields (ordered by official documentation) + background: Optional[bool] = False + include: Optional[ + List[ + Literal[ + "code_interpreter_call.outputs", + "computer_call_output.output.image_url", + "file_search_call.results", + "message.input_image.image_url", + "message.output_text.logprobs", + "reasoning.encrypted_content", + ] + ] + ] = None + input: Union[str, List[ResponseInputOutputItem]] + instructions: Optional[str] = None + max_output_tokens: Optional[int] = None + max_tool_calls: Optional[int] = None + metadata: Optional[Dict[str, Any]] = None + model: Optional[str] = None # Made optional to match vLLM + parallel_tool_calls: Optional[bool] = True + previous_response_id: Optional[str] = None + reasoning: Optional[ResponseReasoningParam] = None + service_tier: Literal["auto", "default", "flex", "scale", "priority"] = "auto" + store: Optional[bool] = True + stream: Optional[bool] = False + temperature: Optional[float] = None + tool_choice: Literal["auto", "required", "none"] = "auto" + tools: List[ResponseTool] = Field(default_factory=list) + top_logprobs: Optional[int] = 0 + top_p: Optional[float] = None + truncation: Optional[Literal["auto", "disabled"]] = "disabled" + user: Optional[str] = None + + # Extra SGLang parameters + request_id: str = Field( + default_factory=lambda: f"resp_{uuid.uuid4().hex}", + description="The request_id related to this request. If the caller does not set it, a random uuid will be generated.", + ) + priority: int = Field(default=0, description="Request priority") + + # SGLang-specific sampling parameters + frequency_penalty: float = 0.0 + presence_penalty: float = 0.0 + stop: Optional[Union[str, List[str]]] = None + top_k: int = -1 + min_p: float = 0.0 + repetition_penalty: float = 1.0 + + # Default sampling parameters + _DEFAULT_SAMPLING_PARAMS = { + "temperature": 0.7, + "top_p": 1.0, + "top_k": -1, + "min_p": 0.0, + "repetition_penalty": 1.0, + } + + def to_sampling_params( + self, default_max_tokens: int, default_params: Optional[Dict] = None + ) -> Dict[str, Any]: + """Convert to sampling parameters for generation.""" + if default_params is None: + default_params = {} + + # Use max_output_tokens if available, otherwise use max_tokens for backwards compatibility + if self.max_output_tokens is not None: + max_tokens = min(self.max_output_tokens, default_max_tokens) + else: + max_tokens = default_max_tokens + + # Avoid exceed the context length by minus 2 token + max_tokens -= 2 + + # Get parameters with defaults + temperature = self.temperature + if temperature is None: + temperature = default_params.get( + "temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"] + ) + + top_p = self.top_p + if top_p is None: + top_p = default_params.get("top_p", self._DEFAULT_SAMPLING_PARAMS["top_p"]) + + params = { + "max_new_tokens": max_tokens, + "temperature": temperature, + "top_p": top_p, + "frequency_penalty": self.frequency_penalty, + "presence_penalty": self.presence_penalty, + "stop": self.stop, + "top_k": self.top_k, + "min_p": self.min_p, + "repetition_penalty": self.repetition_penalty, + } + + # Apply any additional default parameters + for key, value in default_params.items(): + if key not in params or params[key] is None: + params[key] = value + + return params + + +class PromptTokenUsageInfo(BaseModel): + """Prompt token usage details.""" + + cached_tokens: int = 0 + + +class ResponsesResponse(BaseModel): + """Response body for v1/responses endpoint.""" + + id: str = Field(default_factory=lambda: f"resp_{time.time()}") + object: Literal["response"] = "response" + created_at: int = Field(default_factory=lambda: int(time.time())) + model: str + + output: List[ + Union[ResponseOutputItem, ResponseReasoningItem, ResponseFunctionToolCall] + ] = Field(default_factory=list) + status: Literal["queued", "in_progress", "completed", "failed", "cancelled"] + usage: Optional[UsageInfo] = None + parallel_tool_calls: bool = True + tool_choice: str = "auto" + tools: List[ResponseTool] = Field(default_factory=list) + + @classmethod + def from_request( + cls, + request: ResponsesRequest, + sampling_params: Any, + model_name: str, + created_time: int, + output: List[ + Union[ResponseOutputItem, ResponseReasoningItem, ResponseFunctionToolCall] + ], + status: str, + usage: Optional[UsageInfo], + ) -> "ResponsesResponse": + """Create a response from a request.""" + return cls( + id=request.request_id, + created_at=created_time, + model=model_name, + output=output, + status=status, + usage=usage, + parallel_tool_calls=request.parallel_tool_calls or True, + tool_choice=request.tool_choice, + tools=request.tools, + ) + + +class RequestResponseMetadata(BaseModel): + """Metadata for request/response tracking.""" + + request_id: str + final_usage_info: Optional[UsageInfo] = None + + +@dataclass +class MessageProcessingResult: + """Result of processing chat messages and applying templates. + + This dataclass encapsulates all the outputs from message processing including + prompt generation, multimodal data extraction, and constraint preparation. + Used internally by OpenAIServingChat to pass processed data between methods. + + Args: + prompt: The final text prompt after applying chat template + prompt_ids: Either the text prompt (str) or tokenized IDs (List[int]) + image_data: Extracted image data from messages, if any + audio_data: Extracted audio data from messages, if any + modalities: List of modality types present in the messages + stop: Combined stop strings from template and request + tool_call_constraint: Optional constraint for structured tool calls + """ + + prompt: str + prompt_ids: Union[str, List[int]] + image_data: Optional[Any] + audio_data: Optional[Any] + video_data: Optional[Any] + modalities: List[str] + stop: List[str] + tool_call_constraint: Optional[Any] = None + + +class ResponseReasoningTextContent(BaseModel): + text: str + type: Literal["reasoning_text"] = "reasoning_text" + + +ResponseInputOutputItem: TypeAlias = Union[ + ResponseInputItemParam, "ResponseReasoningItem", ResponseFunctionToolCall +] diff --git a/python/sglang/srt/entrypoints/openai/serving_base.py b/python/sglang/srt/entrypoints/openai/serving_base.py new file mode 100644 index 000000000..28b317e6d --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/serving_base.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import json +import logging +import uuid +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Optional, Union + +from fastapi import HTTPException, Request +from fastapi.responses import ORJSONResponse, StreamingResponse + +from sglang.srt.entrypoints.openai.protocol import ErrorResponse, OpenAIServingRequest +from sglang.srt.managers.io_struct import GenerateReqInput + +if TYPE_CHECKING: + from sglang.srt.managers.tokenizer_manager import TokenizerManager + +logger = logging.getLogger(__name__) + + +# Base class for specific endpoint handlers +class OpenAIServingBase(ABC): + """Abstract base class for OpenAI endpoint handlers""" + + def __init__(self, tokenizer_manager: TokenizerManager): + self.tokenizer_manager = tokenizer_manager + + async def handle_request( + self, request: OpenAIServingRequest, raw_request: Request + ) -> Union[Any, StreamingResponse, ErrorResponse]: + """Handle the specific request type with common pattern""" + try: + # Validate request + error_msg = self._validate_request(request) + if error_msg: + return self.create_error_response(error_msg) + + # Convert to internal format + adapted_request, processed_request = self._convert_to_internal_request( + request + ) + + # Note(Xinyuan): raw_request below is only used for detecting the connection of the client + if hasattr(request, "stream") and request.stream: + return await self._handle_streaming_request( + adapted_request, processed_request, raw_request + ) + else: + return await self._handle_non_streaming_request( + adapted_request, processed_request, raw_request + ) + except HTTPException as e: + return self.create_error_response( + message=e.detail, err_type=str(e.status_code), status_code=e.status_code + ) + except Exception as e: + logger.exception(f"Error in request: {e}") + return self.create_error_response( + message=f"Internal server error: {str(e)}", + err_type="InternalServerError", + status_code=500, + ) + + @abstractmethod + def _request_id_prefix(self) -> str: + """Generate request ID based on request type""" + pass + + def _generate_request_id_base(self, request: OpenAIServingRequest) -> Optional[str]: + """Generate request ID based on request type""" + return None + + # TODO(chang): the rid is used in io_strcut check and often violates `The rid should be a list` AssertionError + # Temporarily return None in this function until the rid logic is clear. + if rid := getattr(request, "rid", None): + return rid + + return f"{self._request_id_prefix()}{uuid.uuid4().hex}" + + @abstractmethod + def _convert_to_internal_request( + self, + request: OpenAIServingRequest, + ) -> tuple[GenerateReqInput, OpenAIServingRequest]: + """Convert OpenAI request to internal format""" + pass + + async def _handle_streaming_request( + self, + adapted_request: GenerateReqInput, + request: OpenAIServingRequest, + raw_request: Request, + ) -> Union[StreamingResponse, ErrorResponse, ORJSONResponse]: + """Handle streaming request + + Override this method in child classes that support streaming requests. + """ + return self.create_error_response( + message=f"{self.__class__.__name__} does not support streaming requests", + err_type="NotImplementedError", + status_code=501, + ) + + async def _handle_non_streaming_request( + self, + adapted_request: GenerateReqInput, + request: OpenAIServingRequest, + raw_request: Request, + ) -> Union[Any, ErrorResponse, ORJSONResponse]: + """Handle non-streaming request + + Override this method in child classes that support non-streaming requests. + """ + return self.create_error_response( + message=f"{self.__class__.__name__} does not support non-streaming requests", + err_type="NotImplementedError", + status_code=501, + ) + + def _validate_request(self, _: OpenAIServingRequest) -> Optional[str]: + """Validate request""" + pass + + def create_error_response( + self, + message: str, + err_type: str = "BadRequestError", + status_code: int = 400, + param: Optional[str] = None, + ) -> ORJSONResponse: + """Create an error response""" + # TODO: remove fastapi dependency in openai and move response handling to the entrypoint + error = ErrorResponse( + object="error", + message=message, + type=err_type, + param=param, + code=status_code, + ) + return ORJSONResponse(content=error.model_dump(), status_code=status_code) + + def create_streaming_error_response( + self, + message: str, + err_type: str = "BadRequestError", + status_code: int = 400, + ) -> str: + """Create a streaming error response""" + error = ErrorResponse( + object="error", + message=message, + type=err_type, + param=None, + code=status_code, + ) + return json.dumps({"error": error.model_dump()}) diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py new file mode 100644 index 000000000..d67cbfde3 --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -0,0 +1,1083 @@ +from __future__ import annotations + +import copy +import json +import logging +import time +import uuid +from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union + +from fastapi import Request +from fastapi.responses import ORJSONResponse, StreamingResponse + +from sglang.srt.entrypoints.openai.protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, + ChatCompletionTokenLogprob, + ChatMessage, + ChoiceLogprobs, + DeltaMessage, + ErrorResponse, + FunctionResponse, + LogProbs, + MessageProcessingResult, + ToolCall, + TopLogprob, +) +from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase +from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor +from sglang.srt.entrypoints.openai.utils import ( + process_hidden_states_from_ret, + to_openai_style_logprobs, +) +from sglang.srt.function_call.function_call_parser import FunctionCallParser +from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.parser.conversation import generate_chat_conv +from sglang.srt.parser.jinja_template_utils import process_content_for_template_format +from sglang.srt.parser.reasoning_parser import ReasoningParser +from sglang.utils import convert_json_schema_to_str + +if TYPE_CHECKING: + from sglang.srt.managers.template_manager import TemplateManager + from sglang.srt.managers.tokenizer_manager import TokenizerManager + +logger = logging.getLogger(__name__) + + +class OpenAIServingChat(OpenAIServingBase): + """Handler for /v1/chat/completions requests""" + + def __init__( + self, + tokenizer_manager: TokenizerManager, + template_manager: TemplateManager, + ): + super().__init__(tokenizer_manager) + self.template_manager = template_manager + self.tool_call_parser = self.tokenizer_manager.server_args.tool_call_parser + + def _request_id_prefix(self) -> str: + return "chatcmpl-" + + def _validate_request(self, request: ChatCompletionRequest) -> Optional[str]: + """Validate that the input is valid.""" + if not request.messages: + return "Messages cannot be empty." + + if ( + isinstance(request.tool_choice, str) + and request.tool_choice.lower() == "required" + and not request.tools + ): + return "Tools cannot be empty if tool choice is set to required." + + max_output_tokens = request.max_completion_tokens or request.max_tokens + server_context_length = self.tokenizer_manager.server_args.context_length + if ( + max_output_tokens + and server_context_length + and max_output_tokens > server_context_length + ): + return ( + f"max_completion_tokens is too large: {max_output_tokens}." + f"This model supports at most {server_context_length} completion tokens." + ) + + if request.response_format and request.response_format.type == "json_schema": + schema = getattr(request.response_format.json_schema, "schema_", None) + if schema is None: + return "schema_ is required for json_schema response format request." + + return None + + def _convert_to_internal_request( + self, + request: ChatCompletionRequest, + ) -> tuple[GenerateReqInput, ChatCompletionRequest]: + reasoning_effort = ( + request.chat_template_kwargs.pop("reasoning_effort", None) + if request.chat_template_kwargs + else None + ) + if reasoning_effort is not None: + request.reasoning_effort = reasoning_effort + + """Convert OpenAI chat completion request to internal format""" + is_multimodal = self.tokenizer_manager.model_config.is_multimodal + + # Process messages and apply chat template + processed_messages = self._process_messages(request, is_multimodal) + + # Build sampling parameters + sampling_params = self._build_sampling_params( + request, + processed_messages.stop, + processed_messages.tool_call_constraint, + ) + + # Handle single vs multiple requests + if is_multimodal: + prompt_kwargs = {"text": processed_messages.prompt} + else: + if isinstance(processed_messages.prompt_ids, str): + prompt_kwargs = {"text": processed_messages.prompt_ids} + else: + prompt_kwargs = {"input_ids": processed_messages.prompt_ids} + + adapted_request = GenerateReqInput( + **prompt_kwargs, + image_data=processed_messages.image_data, + video_data=processed_messages.video_data, + audio_data=processed_messages.audio_data, + sampling_params=sampling_params, + return_logprob=request.logprobs, + logprob_start_len=-1, + top_logprobs_num=request.top_logprobs or 0, + stream=request.stream, + return_text_in_logprobs=True, + modalities=processed_messages.modalities, + lora_path=request.lora_path, + bootstrap_host=request.bootstrap_host, + bootstrap_port=request.bootstrap_port, + bootstrap_room=request.bootstrap_room, + return_hidden_states=request.return_hidden_states, + rid=request.rid, + ) + + return adapted_request, request + + def _process_messages( + self, request: ChatCompletionRequest, is_multimodal: bool + ) -> MessageProcessingResult: + """Process chat messages and apply chat template""" + is_gpt_oss = ( + hasattr(self.tokenizer_manager.model_config, "hf_config") + and hasattr(self.tokenizer_manager.model_config.hf_config, "model_type") + and self.tokenizer_manager.model_config.hf_config.model_type == "gpt_oss" + ) + + # GptOss model needs to keep special tokens for harmony parsing + if is_gpt_oss: + request.skip_special_tokens = False + + tool_call_constraint = None + + # Apply chat template and its stop strings + tools = None + if request.tools and request.tool_choice != "none": + request.skip_special_tokens = False + if not isinstance(request.tool_choice, str): + tools = [ + item.function.model_dump() + for item in request.tools + if item.function.name == request.tool_choice.function.name + ] + else: + tools = [item.function.model_dump() for item in request.tools] + if self.tool_call_parser: + parser = FunctionCallParser(request.tools, self.tool_call_parser) + tool_call_constraint = parser.get_structure_constraint( + request.tool_choice + ) + + # Use chat template + if self.template_manager.chat_template_name is None: + result = self._apply_jinja_template(request, tools, is_multimodal) + else: + result = self._apply_conversation_template(request, is_multimodal) + + result.tool_call_constraint = tool_call_constraint + return result + + def _apply_jinja_template( + self, + request: ChatCompletionRequest, + tools: Optional[List[Dict]], + is_multimodal: bool, + ) -> MessageProcessingResult: + """Apply Jinja chat template""" + prompt = "" + prompt_ids = [] + openai_compatible_messages = [] + image_data = [] + video_data = [] + audio_data = [] + modalities = [] + + template_content_format = self.template_manager.jinja_template_content_format + + for message in request.messages: + if message.content is None: + message.content = "" + msg_dict = message.model_dump() + + # Process content based on detected template format + processed_msg = process_content_for_template_format( + msg_dict, + template_content_format, + image_data, + video_data, + audio_data, + modalities, + ) + + # per the Transformers docs & maintainers, tool call arguments in + # assistant-role messages with tool_calls need to be dicts not JSON str - + # this is how tool-use chat templates will expect them moving forwards + # so, for messages that have tool_calls, parse the string (which we get + # from openAI format) to dict + if ( + processed_msg["role"] == "assistant" + and "tool_calls" in processed_msg + and isinstance(processed_msg["tool_calls"], list) + ): + for item in processed_msg["tool_calls"]: + if "arguments" in item["function"] and isinstance( + item["function"]["arguments"], str + ): + item["function"]["arguments"] = json.loads( + item["function"]["arguments"] + ) + + openai_compatible_messages.append(processed_msg) + + # Handle assistant prefix for continue_final_message + assistant_prefix = None + if ( + openai_compatible_messages + and openai_compatible_messages[-1]["role"] == "assistant" + ): + if request.continue_final_message: + assistant_prefix = openai_compatible_messages[-1]["content"] + openai_compatible_messages = openai_compatible_messages[:-1] + + try: + prompt_ids = self.tokenizer_manager.tokenizer.apply_chat_template( + openai_compatible_messages, + tokenize=True, + add_generation_prompt=True, + tools=tools, + reasoning_effort=request.reasoning_effort, + **( + request.chat_template_kwargs if request.chat_template_kwargs else {} + ), + ) + except Exception: + # This except branch will be triggered when the chosen model + # has a different tools input format that is not compatible + # with openAI's apply_chat_template tool_call format, like Mistral. + tools = ( + [t if "function" in t else {"function": t} for t in tools] + if tools + else None + ) + prompt_ids = self.tokenizer_manager.tokenizer.apply_chat_template( + openai_compatible_messages, + tokenize=True, + add_generation_prompt=True, + tools=tools, + reasoning_effort=request.reasoning_effort, + **( + request.chat_template_kwargs if request.chat_template_kwargs else {} + ), + ) + + if assistant_prefix: + encoded = self.tokenizer_manager.tokenizer.encode(assistant_prefix) + if encoded and encoded[0] == self.tokenizer_manager.tokenizer.bos_token_id: + encoded = encoded[1:] + prompt_ids += encoded + + if is_multimodal: + prompt = self.tokenizer_manager.tokenizer.decode(prompt_ids) + + stop = request.stop + image_data = image_data if image_data else None + audio_data = audio_data if audio_data else None + video_data = video_data if video_data else None + modalities = modalities if modalities else [] + return MessageProcessingResult( + prompt=prompt, + prompt_ids=prompt_ids, + image_data=image_data, + video_data=video_data, + audio_data=audio_data, + modalities=modalities, + stop=stop, + ) + + def _apply_conversation_template( + self, + request: ChatCompletionRequest, + is_multimodal: bool, + ) -> MessageProcessingResult: + """Apply conversation template""" + prompt = "" + prompt_ids = [] + conv = generate_chat_conv(request, self.template_manager.chat_template_name) + + # If we should continue the final assistant message, adjust the conversation. + if ( + request.continue_final_message + and request.messages + and request.messages[-1].role == "assistant" + ): + # Remove the auto-added blank assistant turn, if present. + if conv.messages and conv.messages[-1][1] is None: + conv.messages.pop() + # Rebuild the prompt from the conversation. + prompt = conv.get_prompt() + # Strip trailing stop tokens or separators that indicate end-of-assistant. + if isinstance(conv.stop_str, list): + for stop_token in conv.stop_str: + if prompt.endswith(stop_token): + prompt = prompt[: -len(stop_token)] + elif isinstance(conv.stop_str, str) and prompt.endswith(conv.stop_str): + prompt = prompt[: -len(conv.stop_str)] + if conv.sep and prompt.endswith(conv.sep): + prompt = prompt[: -len(conv.sep)] + if getattr(conv, "sep2", None) and prompt.endswith(conv.sep2): + prompt = prompt[: -len(conv.sep2)] + else: + prompt = conv.get_prompt() + if self._get_enable_thinking_from_request(request): + prompt += "" # Note(Xinyuan): hard code thinking token + + image_data = conv.image_data if conv.image_data else None + video_data = conv.video_data if conv.video_data else None + audio_data = conv.audio_data if conv.audio_data else None + modalities = conv.modalities if conv.modalities else [] + stop = copy.copy(conv.stop_str or [] if not request.ignore_eos else []) + + if request.stop: + if isinstance(request.stop, str): + stop.append(request.stop) + else: + stop.extend(request.stop) + + if not is_multimodal: + prompt_ids = self.tokenizer_manager.tokenizer.encode(prompt) + + return MessageProcessingResult( + prompt=prompt, + prompt_ids=prompt_ids, + image_data=image_data, + video_data=video_data, + audio_data=audio_data, + modalities=modalities, + stop=stop, + ) + + def _build_sampling_params( + self, + request: ChatCompletionRequest, + stop: List[str], + tool_call_constraint: Optional[Any], + ) -> Dict[str, Any]: + """Build sampling parameters for the request""" + + sampling_params = { + "temperature": request.temperature, + "max_new_tokens": request.max_tokens or request.max_completion_tokens, + "min_new_tokens": request.min_tokens, + "stop": stop, + "stop_token_ids": request.stop_token_ids, + "top_p": request.top_p, + "top_k": request.top_k, + "min_p": request.min_p, + "presence_penalty": request.presence_penalty, + "frequency_penalty": request.frequency_penalty, + "repetition_penalty": request.repetition_penalty, + "regex": request.regex, + "ebnf": request.ebnf, + "n": request.n, + "no_stop_trim": request.no_stop_trim, + "ignore_eos": request.ignore_eos, + "skip_special_tokens": request.skip_special_tokens, + "logit_bias": request.logit_bias, + } + + if request.response_format and request.response_format.type == "json_schema": + sampling_params["json_schema"] = convert_json_schema_to_str( + request.response_format.json_schema.schema_ + ) + elif request.response_format and request.response_format.type == "json_object": + sampling_params["json_schema"] = '{"type": "object"}' + elif ( + request.response_format and request.response_format.type == "structural_tag" + ): + sampling_params["structural_tag"] = convert_json_schema_to_str( + request.response_format.model_dump(by_alias=True) + ) + + # Check if there are already existing output constraints + has_existing_constraints = ( + sampling_params.get("regex") + or sampling_params.get("ebnf") + or sampling_params.get("structural_tag") + or sampling_params.get("json_schema") + ) + + if tool_call_constraint and has_existing_constraints: + logger.warning("Constrained decoding is not compatible with tool calls.") + elif tool_call_constraint: + constraint_type, constraint_value = tool_call_constraint + if constraint_type == "structural_tag": + sampling_params[constraint_type] = convert_json_schema_to_str( + constraint_value.model_dump(by_alias=True) + ) + else: + sampling_params[constraint_type] = constraint_value + return sampling_params + + async def _handle_streaming_request( + self, + adapted_request: GenerateReqInput, + request: ChatCompletionRequest, + raw_request: Request, + ) -> StreamingResponse: + """Handle streaming chat completion request""" + return StreamingResponse( + self._generate_chat_stream(adapted_request, request, raw_request), + media_type="text/event-stream", + background=self.tokenizer_manager.create_abort_task(adapted_request), + ) + + async def _generate_chat_stream( + self, + adapted_request: GenerateReqInput, + request: ChatCompletionRequest, + raw_request: Request, + ) -> AsyncGenerator[str, None]: + """Generate streaming chat completion response""" + # Parsers for tool calls and reasoning + parser_dict = {} + reasoning_parser_dict = {} + + # State tracking for streaming + is_firsts = {} + stream_buffers = {} + n_prev_tokens = {} + has_tool_calls = {} + finish_reasons = {} + + # Usage tracking + prompt_tokens = {} + completion_tokens = {} + cached_tokens = {} + hidden_states = {} + + try: + async for content in self.tokenizer_manager.generate_request( + adapted_request, raw_request + ): + index = content.get("index", 0) + + prompt_tokens[index] = content["meta_info"]["prompt_tokens"] + completion_tokens[index] = content["meta_info"]["completion_tokens"] + cached_tokens[index] = content["meta_info"].get("cached_tokens", 0) + hidden_states[index] = content["meta_info"].get("hidden_states", None) + + # Handle logprobs + choice_logprobs = None + if request.logprobs: + choice_logprobs = self._process_streaming_logprobs( + content, n_prev_tokens.get(index, 0) + ) + n_prev_tokens[index] = len( + content["meta_info"]["output_token_logprobs"] + ) + + finish_reason = content["meta_info"]["finish_reason"] + finish_reason_type = finish_reason["type"] if finish_reason else None + + # Track finish_reason for each index + if finish_reason_type: + finish_reasons[index] = finish_reason + + # First chunk with role + if is_firsts.get(index, True): + is_firsts[index] = False + delta = DeltaMessage(role="assistant", content="") + choice_data = ChatCompletionResponseStreamChoice( + index=index, + delta=delta, + finish_reason=None, + logprobs=None, + ) + chunk = ChatCompletionStreamResponse( + id=content["meta_info"]["id"], + created=int(time.time()), + choices=[choice_data], + model=request.model, + ) + yield f"data: {chunk.model_dump_json()}\n\n" + + stream_buffer = stream_buffers.get(index, "") + delta = content["text"][len(stream_buffer) :] + stream_buffers[index] = stream_buffer + delta + + # Handle reasoning content + if ( + self.tokenizer_manager.server_args.reasoning_parser + and request.separate_reasoning + ): + reasoning_text, delta = self._process_reasoning_stream( + index, delta, reasoning_parser_dict, content, request + ) + if reasoning_text: + choice_data = ChatCompletionResponseStreamChoice( + index=index, + delta=DeltaMessage(reasoning_content=reasoning_text), + finish_reason=None, + ) + chunk = ChatCompletionStreamResponse( + id=content["meta_info"]["id"], + created=int(time.time()), + choices=[choice_data], + model=request.model, + ) + yield f"data: {chunk.model_dump_json()}\n\n" + + # Handle tool calls + if ( + request.tool_choice != "none" + and request.tools + and self.tool_call_parser + ): + async for chunk in self._process_tool_call_stream( + index, + delta, + parser_dict, + content, + request, + has_tool_calls, + ): + if chunk: + yield chunk + + # Send any remaining tool call arguments when generation finishes + if finish_reason_type is not None and index in parser_dict: + parser = parser_dict[index] + remaining_chunk = self._check_for_unstreamed_tool_args( + parser, content, request, index + ) + if remaining_chunk: + yield remaining_chunk + + else: + # Regular content + if delta: + choice_data = ChatCompletionResponseStreamChoice( + index=index, + delta=DeltaMessage(content=delta), + finish_reason=None, + matched_stop=None, + logprobs=choice_logprobs, + ) + chunk = ChatCompletionStreamResponse( + id=content["meta_info"]["id"], + created=int(time.time()), + choices=[choice_data], + model=request.model, + ) + yield f"data: {chunk.model_dump_json()}\n\n" + + # Send finish_reason chunks for each index that completed + for idx, finish_reason_data in finish_reasons.items(): + finish_reason_type = finish_reason_data["type"] + + # Change finish_reason to "tool_calls" if we had tool calls and stopped naturally + final_finish_reason = finish_reason_type + if has_tool_calls.get(idx, False) and finish_reason_type == "stop": + final_finish_reason = "tool_calls" + + finish_reason_chunk = ChatCompletionStreamResponse( + id=content["meta_info"][ + "id" + ], # NOTE: openai uses the same chatcmpl-id for all indices + created=int(time.time()), + choices=[ + ChatCompletionResponseStreamChoice( + index=idx, + delta=DeltaMessage(), + finish_reason=final_finish_reason, + matched_stop=( + finish_reason_data["matched"] + if "matched" in finish_reason_data + else None + ), + ) + ], + model=request.model, + usage=None, + ) + yield f"data: {finish_reason_chunk.model_dump_json()}\n\n" + + # Send hidden states if requested + if request.return_hidden_states and hidden_states: + for index, choice_hidden_states in hidden_states.items(): + if choice_hidden_states: + last_token_hidden_states = ( + choice_hidden_states[-1] + if len(choice_hidden_states) > 1 + else [] + ) + hidden_states_chunk = ChatCompletionStreamResponse( + id=content["meta_info"]["id"], + created=int(time.time()), + choices=[ + ChatCompletionResponseStreamChoice( + index=index, + delta=DeltaMessage( + hidden_states=last_token_hidden_states + ), + finish_reason=None, # Hidden states don't need finish_reason + ) + ], + model=request.model, + ) + yield f"data: {hidden_states_chunk.model_dump_json()}\n\n" + + # Additional usage chunk + if request.stream_options and request.stream_options.include_usage: + usage = UsageProcessor.calculate_streaming_usage( + prompt_tokens, + completion_tokens, + cached_tokens, + n_choices=request.n, + enable_cache_report=self.tokenizer_manager.server_args.enable_cache_report, + ) + usage_chunk = ChatCompletionStreamResponse( + id=content["meta_info"]["id"], + created=int(time.time()), + choices=[], # Empty choices array as per OpenAI spec + model=request.model, + usage=usage, + ) + yield f"data: {usage_chunk.model_dump_json()}\n\n" + + except ValueError as e: + error = self.create_streaming_error_response(str(e)) + yield f"data: {error}\n\n" + + yield "data: [DONE]\n\n" + + async def _handle_non_streaming_request( + self, + adapted_request: GenerateReqInput, + request: ChatCompletionRequest, + raw_request: Request, + ) -> Union[ChatCompletionResponse, ErrorResponse, ORJSONResponse]: + """Handle non-streaming chat completion request""" + try: + ret = await self.tokenizer_manager.generate_request( + adapted_request, raw_request + ).__anext__() + except ValueError as e: + return self.create_error_response(str(e)) + + if not isinstance(ret, list): + ret = [ret] + + response = self._build_chat_response( + request, + ret, + int(time.time()), + ) + + return response + + def _build_chat_response( + self, + request: ChatCompletionRequest, + ret: List[Dict[str, Any]], + created: int, + ) -> Union[ChatCompletionResponse, ORJSONResponse]: + """Build chat completion response from generation results""" + choices = [] + + for idx, ret_item in enumerate(ret): + # Process logprobs + choice_logprobs = None + if request.logprobs: + choice_logprobs = self._process_response_logprobs(ret_item) + + # Handle hidden states + hidden_states = process_hidden_states_from_ret(ret_item, request) + + finish_reason = ret_item["meta_info"]["finish_reason"] + text = ret_item["text"] + + # Handle reasoning content + reasoning_text = None + reasoning_parser = self.tokenizer_manager.server_args.reasoning_parser + if reasoning_parser and request.separate_reasoning: + is_force_reasoning = ( + self.template_manager.force_reasoning + or self._get_enable_thinking_from_request(request) + ) + try: + parser = ReasoningParser( + model_type=reasoning_parser, + stream_reasoning=False, + force_reasoning=is_force_reasoning, + ) + reasoning_text, text = parser.parse_non_stream(text) + except Exception as e: + logger.error(f"Reasoning parsing error: {e}") + return self.create_error_response( + "Failed to parse reasoning content", + err_type="InternalServerError", + status_code=500, + ) + + # Handle tool calls + tool_calls = None + if ( + request.tool_choice != "none" + and request.tools + and self.tool_call_parser + ): + tool_calls, text, finish_reason = self._process_tool_calls( + text, request.tools, finish_reason + ) + + choice_data = ChatCompletionResponseChoice( + index=idx, + message=ChatMessage( + role="assistant", + content=text if text else None, + tool_calls=tool_calls, + reasoning_content=reasoning_text if reasoning_text else None, + ), + logprobs=choice_logprobs, + finish_reason=finish_reason["type"] if finish_reason else None, + matched_stop=( + finish_reason["matched"] + if finish_reason and "matched" in finish_reason + else None + ), + hidden_states=hidden_states, + ) + choices.append(choice_data) + + # Calculate usage + usage = UsageProcessor.calculate_response_usage( + ret, + n_choices=request.n, + enable_cache_report=self.tokenizer_manager.server_args.enable_cache_report, + ) + + return ChatCompletionResponse( + id=ret[0]["meta_info"]["id"], + created=created, + model=request.model, + choices=choices, + usage=usage, + metadata={"weight_version": ret[0]["meta_info"]["weight_version"]}, + ) + + def _process_logprobs_tokens( + self, logprobs: LogProbs, use_token_index: bool = False + ) -> List[ChatCompletionTokenLogprob]: + """Common helper to process logprobs tokens for both streaming and non-streaming + + Args: + logprobs: LogProbs data from model + use_token_index: True for non-streaming (use token_idx), False for streaming (use index 0) + """ + token_logprobs = [] + + for token_idx, (token, logprob) in enumerate( + zip(logprobs.tokens, logprobs.token_logprobs) + ): + token_bytes = list(token.encode("utf-8")) + top_logprobs = [] + if logprobs.top_logprobs: + # - Non-streaming (use_token_index=True): uses token_idx for full data + # - Streaming (use_token_index=False): uses index 0 for pre-sliced data + top_logprobs_idx = token_idx if use_token_index else 0 + for top_token, top_logprob in logprobs.top_logprobs[ + top_logprobs_idx + ].items(): + top_token_bytes = list(top_token.encode("utf-8")) + top_logprobs.append( + TopLogprob( + token=top_token, + bytes=top_token_bytes, + logprob=top_logprob, + ) + ) + token_logprobs.append( + ChatCompletionTokenLogprob( + token=token, + bytes=token_bytes, + logprob=logprob, + top_logprobs=top_logprobs, + ) + ) + + return token_logprobs + + def _process_response_logprobs(self, ret_item: Dict[str, Any]) -> ChoiceLogprobs: + """Process logprobs for non-streaming response""" + logprobs = to_openai_style_logprobs( + output_token_logprobs=ret_item["meta_info"]["output_token_logprobs"], + output_top_logprobs=ret_item["meta_info"].get("output_top_logprobs", None), + ) + + token_logprobs = self._process_logprobs_tokens(logprobs, use_token_index=True) + return ChoiceLogprobs(content=token_logprobs) + + def _process_tool_calls( + self, + text: str, + tools: List[Any], + finish_reason: Dict[str, Any], + ) -> tuple[Optional[List[ToolCall]], str, Dict[str, Any]]: + """Process tool calls in the response""" + parser = FunctionCallParser(tools, self.tool_call_parser) + if parser.has_tool_call(text): + if finish_reason["type"] == "stop": + finish_reason["type"] = "tool_calls" + finish_reason["matched"] = None + try: + text, call_info_list = parser.parse_non_stream(text) + tool_calls = [] + for call_info in call_info_list: + # For Kimi-K2, align tool_call_id with the model format: functions.{name}:{index} + if ( + self.tool_call_parser == "kimi_k2" + and call_info.name is not None + ): + tool_id = f"functions.{call_info.name}:{call_info.tool_index}" + else: + tool_id = f"call_{uuid.uuid4().hex[:24]}" + + tool_calls.append( + ToolCall( + id=tool_id, + index=getattr(call_info, "tool_index", None), + function=FunctionResponse( + name=call_info.name, arguments=call_info.parameters + ), + ) + ) + return tool_calls, text, finish_reason + except Exception as e: + logger.error(f"Tool call parsing error: {e}") + # Return error but don't fail the whole request + return None, text, finish_reason + + return None, text, finish_reason + + def _process_streaming_logprobs( + self, content: Dict[str, Any], n_prev_token: int + ) -> ChoiceLogprobs: + """Process logprobs for streaming response""" + logprobs = to_openai_style_logprobs( + output_token_logprobs=content["meta_info"]["output_token_logprobs"][ + n_prev_token: + ], + output_top_logprobs=content["meta_info"].get("output_top_logprobs", [])[ + n_prev_token: + ], + ) + + token_logprobs = self._process_logprobs_tokens(logprobs, use_token_index=False) + return ChoiceLogprobs(content=token_logprobs) + + def _process_reasoning_stream( + self, + index: int, + delta: str, + reasoning_parser_dict: Dict[int, ReasoningParser], + content: Dict[str, Any], + request: ChatCompletionRequest, + ) -> tuple[Optional[str], str]: + """Process reasoning content in streaming response""" + if index not in reasoning_parser_dict: + is_force_reasoning = ( + self.template_manager.force_reasoning + or self._get_enable_thinking_from_request(request) + ) + reasoning_parser_dict[index] = ReasoningParser( + self.tokenizer_manager.server_args.reasoning_parser, + request.stream_reasoning, + is_force_reasoning, + ) + reasoning_parser = reasoning_parser_dict[index] + return reasoning_parser.parse_stream_chunk(delta) + + def _get_enable_thinking_from_request(self, request: ChatCompletionRequest) -> bool: + """Extracts the 'enable_thinking' flag from request chat_template_kwargs. + + NOTE: This parameter is only useful for models that support enable_thinking + flag, such as Qwen3. + + Args: + request_obj: The request object (or an item from a list of requests). + Returns: + The boolean value of 'enable_thinking' if found, otherwise False. + """ + if hasattr(request, "chat_template_kwargs") and request.chat_template_kwargs: + # For Qwen3 models, `enable_thinking` is supported. + if request.chat_template_kwargs.get("enable_thinking") is not None: + return request.chat_template_kwargs.get("enable_thinking") + # For DeepSeek-V3.1 models, `thinking` is supported. + elif request.chat_template_kwargs.get("thinking") is not None: + return request.chat_template_kwargs.get("thinking") + else: + return False + return False + + async def _process_tool_call_stream( + self, + index: int, + delta: str, + parser_dict: Dict[int, FunctionCallParser], + content: Dict[str, Any], + request: ChatCompletionRequest, + has_tool_calls: Dict[int, bool], + ): + """Process tool calls in streaming response""" + if index not in parser_dict: + parser_dict[index] = FunctionCallParser( + tools=request.tools, + tool_call_parser=self.tool_call_parser, + ) + parser = parser_dict[index] + + normal_text, calls = parser.parse_stream_chunk(delta) + + # Yield normal text + if normal_text: + choice_data = ChatCompletionResponseStreamChoice( + index=index, + delta=DeltaMessage(content=normal_text), + finish_reason=None, + ) + chunk = ChatCompletionStreamResponse( + id=content["meta_info"]["id"], + created=int(time.time()), + choices=[choice_data], + model=request.model, + ) + yield f"data: {chunk.model_dump_json()}\n\n" + + # Yield tool calls + for call_item in calls: + # Mark that this choice has tool calls + has_tool_calls[index] = True + + # Tool call ID should be generated only once per tool call + if call_item.name: + # First chunk: include ID and function name + if self.tool_call_parser == "kimi_k2": + # Align with Kimi-K2 format: functions.{name}:{index} + tool_call_id = f"functions.{call_item.name}:{call_item.tool_index}" + else: + tool_call_id = f"call_{uuid.uuid4().hex[:24]}" + function_name = call_item.name + else: + # Subsequent chunks: null ID and name for argument deltas + tool_call_id = None + function_name = None + + tool_call = ToolCall( + id=tool_call_id, + index=call_item.tool_index, + function=FunctionResponse( + name=function_name, + arguments=call_item.parameters, + ), + ) + + choice_data = ChatCompletionResponseStreamChoice( + index=index, + delta=DeltaMessage(tool_calls=[tool_call]), + finish_reason=None, + ) + chunk = ChatCompletionStreamResponse( + id=content["meta_info"]["id"], + created=int(time.time()), + choices=[choice_data], + model=request.model, + ) + yield f"data: {chunk.model_dump_json()}\n\n" + + def _check_for_unstreamed_tool_args( + self, + parser: FunctionCallParser, + content: Dict[str, Any], + request: ChatCompletionRequest, + index: int, + ) -> Optional[str]: + """ + Check for any remaining tool call arguments that need to be streamed + when generation finishes. This ensures tool calls are properly completed + even if the model generates the final arguments in the last chunk. + """ + # Only check if we have tool calls and the parser has tracked data + if ( + not hasattr(parser.detector, "prev_tool_call_arr") + or not parser.detector.prev_tool_call_arr + ): + return None + + if ( + not hasattr(parser.detector, "streamed_args_for_tool") + or not parser.detector.streamed_args_for_tool + ): + return None + + # Get the last tool call that was being processed + tool_index = len(parser.detector.prev_tool_call_arr) - 1 + if tool_index < 0 or tool_index >= len(parser.detector.streamed_args_for_tool): + return None + + # Get expected vs actual arguments + expected_args = parser.detector.prev_tool_call_arr[tool_index].get( + "arguments", {} + ) + expected_call = json.dumps(expected_args, ensure_ascii=False) + actual_call = parser.detector.streamed_args_for_tool[tool_index] + + # Check if there are remaining arguments to send + remaining_call = ( + expected_call.replace(actual_call, "", 1) + if actual_call in expected_call + else "" + ) + + if remaining_call: + # Create tool call chunk with remaining arguments + tool_call = ToolCall( + id=None, # No ID for argument deltas + index=tool_index, + function=FunctionResponse( + name=None, # No name for argument deltas + arguments=remaining_call, + ), + ) + + choice_data = ChatCompletionResponseStreamChoice( + index=index, + delta=DeltaMessage(tool_calls=[tool_call]), + finish_reason=None, # Don't send finish_reason with this chunk + ) + + chunk = ChatCompletionStreamResponse( + id=content["meta_info"]["id"], + created=int(time.time()), + choices=[choice_data], + model=request.model, + ) + + return f"data: {chunk.model_dump_json()}\n\n" + + return None diff --git a/python/sglang/srt/entrypoints/openai/serving_completions.py b/python/sglang/srt/entrypoints/openai/serving_completions.py new file mode 100644 index 000000000..6fe02d325 --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/serving_completions.py @@ -0,0 +1,455 @@ +from __future__ import annotations + +import logging +import time +from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union + +from fastapi import Request +from fastapi.responses import ORJSONResponse, StreamingResponse + +from sglang.srt.entrypoints.openai.protocol import ( + CompletionRequest, + CompletionResponse, + CompletionResponseChoice, + CompletionResponseStreamChoice, + CompletionStreamResponse, + ErrorResponse, +) +from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase +from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor +from sglang.srt.entrypoints.openai.utils import ( + process_hidden_states_from_ret, + to_openai_style_logprobs, +) +from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.parser.code_completion_parser import ( + generate_completion_prompt_from_request, +) +from sglang.utils import convert_json_schema_to_str + +if TYPE_CHECKING: + from sglang.srt.managers.template_manager import TemplateManager + from sglang.srt.managers.tokenizer_manager import TokenizerManager + +logger = logging.getLogger(__name__) + + +class OpenAIServingCompletion(OpenAIServingBase): + """Handler for /v1/completion requests""" + + def __init__( + self, + tokenizer_manager: TokenizerManager, + template_manager: TemplateManager, + ): + super().__init__(tokenizer_manager) + self.template_manager = template_manager + + def _request_id_prefix(self) -> str: + return "cmpl-" + + def _validate_request(self, request: CompletionRequest) -> Optional[str]: + """Validate that the input is valid.""" + prompt = request.prompt + if not prompt or (isinstance(prompt, list) and all(not p for p in prompt)): + return "Prompt cannot be empty" + + return None + + def _convert_to_internal_request( + self, + request: CompletionRequest, + ) -> tuple[GenerateReqInput, CompletionRequest]: + """Convert OpenAI completion request to internal format""" + # NOTE: with openai API, the prompt's logprobs are always not computed + if request.echo and request.logprobs: + logger.warning( + "Echo is not compatible with logprobs. " + "To compute logprobs of input prompt, please use the native /generate API." + ) + # Process prompt + prompt = request.prompt + if self.template_manager.completion_template_name is not None: + prompt = generate_completion_prompt_from_request(request) + + # Set logprob start length based on echo and logprobs + if request.echo and request.logprobs: + logprob_start_len = 0 + else: + logprob_start_len = -1 + + # Build sampling parameters + sampling_params = self._build_sampling_params(request) + + # Determine prompt format + if isinstance(prompt, str) or ( + isinstance(prompt, list) and isinstance(prompt[0], str) + ): + prompt_kwargs = {"text": prompt} + else: + prompt_kwargs = {"input_ids": prompt} + + adapted_request = GenerateReqInput( + **prompt_kwargs, + sampling_params=sampling_params, + return_logprob=request.logprobs is not None, + top_logprobs_num=request.logprobs if request.logprobs is not None else 0, + logprob_start_len=logprob_start_len, + return_text_in_logprobs=True, + stream=request.stream, + lora_path=request.lora_path, + bootstrap_host=request.bootstrap_host, + bootstrap_port=request.bootstrap_port, + bootstrap_room=request.bootstrap_room, + return_hidden_states=request.return_hidden_states, + rid=request.rid, + ) + + return adapted_request, request + + def _build_sampling_params(self, request: CompletionRequest) -> Dict[str, Any]: + """Build sampling parameters for the request""" + # Start with common parameters + sampling_params = { + "temperature": request.temperature, + "max_new_tokens": request.max_tokens, + "min_new_tokens": request.min_tokens, + "stop": request.stop, + "stop_token_ids": request.stop_token_ids, + "top_p": request.top_p, + "top_k": request.top_k, + "min_p": request.min_p, + "presence_penalty": request.presence_penalty, + "frequency_penalty": request.frequency_penalty, + "repetition_penalty": request.repetition_penalty, + "regex": request.regex, + "json_schema": request.json_schema, + "ebnf": request.ebnf, + "n": request.n, + "no_stop_trim": request.no_stop_trim, + "ignore_eos": request.ignore_eos, + "skip_special_tokens": request.skip_special_tokens, + "logit_bias": request.logit_bias, + } + + # Handle response_format constraints + if request.response_format and request.response_format.type == "json_schema": + sampling_params["json_schema"] = convert_json_schema_to_str( + request.response_format.json_schema.schema_ + ) + elif request.response_format and request.response_format.type == "json_object": + sampling_params["json_schema"] = '{"type": "object"}' + elif ( + request.response_format and request.response_format.type == "structural_tag" + ): + sampling_params["structural_tag"] = convert_json_schema_to_str( + request.response_format.model_dump(by_alias=True) + ) + + return sampling_params + + async def _handle_streaming_request( + self, + adapted_request: GenerateReqInput, + request: CompletionRequest, + raw_request: Request, + ) -> StreamingResponse: + """Handle streaming completion request""" + return StreamingResponse( + self._generate_completion_stream(adapted_request, request, raw_request), + media_type="text/event-stream", + background=self.tokenizer_manager.create_abort_task(adapted_request), + ) + + async def _generate_completion_stream( + self, + adapted_request: GenerateReqInput, + request: CompletionRequest, + raw_request: Request, + ) -> AsyncGenerator[str, None]: + """Generate streaming completion response""" + created = int(time.time()) + + # State tracking for streaming + stream_buffers = {} + n_prev_tokens = {} + + # Usage tracking + prompt_tokens = {} + completion_tokens = {} + cached_tokens = {} + hidden_states = {} + + try: + async for content in self.tokenizer_manager.generate_request( + adapted_request, raw_request + ): + index = content.get("index", 0) + + text = content["text"] + prompt_tokens[index] = content["meta_info"]["prompt_tokens"] + completion_tokens[index] = content["meta_info"]["completion_tokens"] + cached_tokens[index] = content["meta_info"].get("cached_tokens", 0) + hidden_states[index] = content["meta_info"].get("hidden_states", None) + + stream_buffer = stream_buffers.get(index, "") + # Handle echo for first chunk + if not stream_buffer: # The first chunk + if request.echo: + echo_text = self._get_echo_text(request, index) + text = echo_text + text + + # Handle logprobs + logprobs = None + if request.logprobs is not None: + # The first chunk and echo is enabled. + if not stream_buffer and request.echo: + input_token_logprobs = content["meta_info"][ + "input_token_logprobs" + ] + input_top_logprobs = content["meta_info"]["input_top_logprobs"] + else: + input_token_logprobs = None + input_top_logprobs = None + + n_prev_token = n_prev_tokens.get(index, 0) + logprobs = to_openai_style_logprobs( + input_token_logprobs=input_token_logprobs, + input_top_logprobs=input_top_logprobs, + output_token_logprobs=content["meta_info"][ + "output_token_logprobs" + ][n_prev_token:], + output_top_logprobs=content["meta_info"]["output_top_logprobs"][ + n_prev_token: + ], + ) + n_prev_tokens[index] = len( + content["meta_info"]["output_token_logprobs"] + ) + + # Generate delta + delta = text[len(stream_buffer) :] + stream_buffers[index] = stream_buffer + delta + finish_reason = content["meta_info"]["finish_reason"] + + choice_data = CompletionResponseStreamChoice( + index=index, + text=delta, + logprobs=logprobs, + finish_reason=finish_reason["type"] if finish_reason else None, + matched_stop=( + finish_reason["matched"] + if finish_reason and "matched" in finish_reason + else None + ), + ) + chunk = CompletionStreamResponse( + id=content["meta_info"]["id"], + created=created, + object="text_completion", + choices=[choice_data], + model=request.model, + ) + + yield f"data: {chunk.model_dump_json()}\n\n" + + if request.return_hidden_states and hidden_states: + for index, choice_hidden_states in hidden_states.items(): + if choice_hidden_states: + last_token_hidden_states = ( + choice_hidden_states[-1] + if len(choice_hidden_states) > 1 + else [] + ) + hidden_states_chunk = CompletionStreamResponse( + id=content["meta_info"]["id"], + created=created, + object="text_completion", + choices=[ + CompletionResponseStreamChoice( + index=index, + text="", + hidden_states=last_token_hidden_states, + finish_reason=None, + ) + ], + model=request.model, + ) + yield f"data: {hidden_states_chunk.model_dump_json()}\n\n" + + # Handle final usage chunk + if request.stream_options and request.stream_options.include_usage: + usage = UsageProcessor.calculate_streaming_usage( + prompt_tokens, + completion_tokens, + cached_tokens, + n_choices=request.n, + enable_cache_report=self.tokenizer_manager.server_args.enable_cache_report, + ) + final_usage_chunk = CompletionStreamResponse( + id=content["meta_info"]["id"], + created=created, + choices=[], + model=request.model, + usage=usage, + ) + final_usage_data = final_usage_chunk.model_dump_json(exclude_none=True) + yield f"data: {final_usage_data}\n\n" + + except Exception as e: + error = self.create_streaming_error_response(str(e)) + yield f"data: {error}\n\n" + + yield "data: [DONE]\n\n" + + async def _handle_non_streaming_request( + self, + adapted_request: GenerateReqInput, + request: CompletionRequest, + raw_request: Request, + ) -> Union[CompletionResponse, ErrorResponse, ORJSONResponse]: + """Handle non-streaming completion request""" + try: + generator = self.tokenizer_manager.generate_request( + adapted_request, raw_request + ) + ret = await generator.__anext__() + except ValueError as e: + return self.create_error_response(str(e)) + + if not isinstance(ret, list): + ret = [ret] + + response = self._build_completion_response( + request, + ret, + int(time.time()), + ) + + return response + + def _build_completion_response( + self, + request: CompletionRequest, + ret: List[Dict[str, Any]], + created: int, + ) -> CompletionResponse: + """Build completion response from generation results""" + choices = [] + echo = False + + # Prepare echo prompts if needed + echo_prompts = [] + if request.echo: + echo_prompts = self._prepare_echo_prompts(request) + echo = True + + for idx, ret_item in enumerate(ret): + text = ret_item["text"] + + # Handle echo + if echo: + prompt_index = idx // request.n + text = echo_prompts[prompt_index] + text + + # Handle logprobs + logprobs = None + if request.logprobs is not None: + if echo: + input_token_logprobs = ret_item["meta_info"]["input_token_logprobs"] + input_top_logprobs = ret_item["meta_info"]["input_top_logprobs"] + else: + input_token_logprobs = None + input_top_logprobs = None + + logprobs = to_openai_style_logprobs( + input_token_logprobs=input_token_logprobs, + input_top_logprobs=input_top_logprobs, + output_token_logprobs=ret_item["meta_info"][ + "output_token_logprobs" + ], + output_top_logprobs=ret_item["meta_info"]["output_top_logprobs"], + ) + + # Handle hidden states + hidden_states = process_hidden_states_from_ret(ret_item, request) + + finish_reason = ret_item["meta_info"]["finish_reason"] + + choice_data = CompletionResponseChoice( + index=idx, + text=text, + logprobs=logprobs, + finish_reason=finish_reason["type"] if finish_reason else None, + matched_stop=( + finish_reason["matched"] + if finish_reason and "matched" in finish_reason + else None + ), + hidden_states=hidden_states, + ) + choices.append(choice_data) + + # Calculate usage + cache_report = self.tokenizer_manager.server_args.enable_cache_report + usage = UsageProcessor.calculate_response_usage( + ret, n_choices=request.n, enable_cache_report=cache_report + ) + + return CompletionResponse( + id=ret[0]["meta_info"]["id"], + model=request.model, + created=created, + choices=choices, + usage=usage, + metadata={"weight_version": ret[0]["meta_info"]["weight_version"]}, + ) + + def _get_echo_text(self, request: CompletionRequest, index: int) -> str: + """Get echo text for streaming response""" + if isinstance(request.prompt, str): + # for the case of single str prompts + return request.prompt + elif isinstance(request.prompt, list): + if isinstance(request.prompt[0], str): + # for the case of multiple str prompts + return request.prompt[index // request.n] + elif isinstance(request.prompt[0], int): + # for the case of single token ids prompt + return self.tokenizer_manager.tokenizer.decode( + request.prompt, skip_special_tokens=True + ) + elif isinstance(request.prompt[0], list) and isinstance( + request.prompt[0][0], int + ): + # for the case of multiple token ids prompts + return self.tokenizer_manager.tokenizer.decode( + request.prompt[index // request.n], + skip_special_tokens=True, + ) + return "" + + def _prepare_echo_prompts(self, request: CompletionRequest) -> List[str]: + """Prepare echo prompts for non-streaming response""" + # TODO: handle the case prompt is token ids + if isinstance(request.prompt, list) and isinstance(request.prompt[0], str): + # for the case of multiple str prompts + return request.prompt + elif isinstance(request.prompt, list) and isinstance(request.prompt[0], list): + # for the case of multiple token ids prompts + return [ + self.tokenizer_manager.tokenizer.decode( + prompt, skip_special_tokens=True + ) + for prompt in request.prompt + ] + elif isinstance(request.prompt, list) and isinstance(request.prompt[0], int): + # for the case of single token ids prompt + return [ + self.tokenizer_manager.tokenizer.decode( + request.prompt, skip_special_tokens=True + ) + ] + else: + # for the case of single str prompt + return [request.prompt] diff --git a/python/sglang/srt/entrypoints/openai/serving_embedding.py b/python/sglang/srt/entrypoints/openai/serving_embedding.py new file mode 100644 index 000000000..63c4fc34a --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/serving_embedding.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from fastapi import Request +from fastapi.responses import ORJSONResponse + +from sglang.srt.entrypoints.openai.protocol import ( + EmbeddingObject, + EmbeddingRequest, + EmbeddingResponse, + ErrorResponse, + MultimodalEmbeddingInput, + UsageInfo, +) +from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase +from sglang.srt.managers.io_struct import EmbeddingReqInput +from sglang.srt.parser.conversation import generate_embedding_convs + +if TYPE_CHECKING: + from sglang.srt.managers.template_manager import TemplateManager + from sglang.srt.managers.tokenizer_manager import TokenizerManager + + +class OpenAIServingEmbedding(OpenAIServingBase): + """Handler for v1/embeddings requests""" + + def __init__( + self, + tokenizer_manager: TokenizerManager, + template_manager: TemplateManager, + ): + super().__init__(tokenizer_manager) + self.template_manager = template_manager + + def _request_id_prefix(self) -> str: + return "embd-" + + def _validate_request(self, request: EmbeddingRequest) -> Optional[str]: + """Validate that the input is not empty or whitespace only.""" + if not (input := request.input): + return "Input cannot be empty" + + # Handle single string + if isinstance(input, str): + if not input.strip(): + return "Input cannot be empty or whitespace only" + return None + + # Handle list inputs + if isinstance(input, list): + if len(input) == 0: + return "Input cannot be empty" + + # Check first element to determine type + first_item = input[0] + + if isinstance(first_item, str): + # List of strings + for i, item in enumerate(input): + if not isinstance(item, str): + return f"All items in input list must be strings" + if not item.strip(): + return f"Input at index {i} cannot be empty or whitespace only" + elif isinstance(first_item, int): + # List of integers (token IDs) + for i, item in enumerate(input): + if not isinstance(item, int): + return f"All items in input list must be integers" + if item < 0: + return f"Token ID at index {i} must be non-negative" + return None + + def _convert_to_internal_request( + self, + request: EmbeddingRequest, + ) -> tuple[EmbeddingReqInput, EmbeddingRequest]: + """Convert OpenAI embedding request to internal format""" + prompt = request.input + + if isinstance(prompt, str): + # Single string input + prompt_kwargs = {"text": prompt} + elif isinstance(prompt, list): + if len(prompt) > 0 and isinstance(prompt[0], str): + prompt_kwargs = {"text": prompt} + elif len(prompt) > 0 and isinstance(prompt[0], MultimodalEmbeddingInput): + # Handle multimodal embedding inputs + texts = [] + images = [] + for item in prompt: + # Use padding for text if None - this could be improved + texts.append(item.text if item.text is not None else "padding") + images.append(item.image if item.image is not None else None) + + generate_prompts = [] + # Check if we have a chat template for multimodal embeddings + if self.template_manager.chat_template_name is not None: + convs = generate_embedding_convs( + texts, images, self.template_manager.chat_template_name + ) + for conv in convs: + generate_prompts.append(conv.get_prompt()) + else: + generate_prompts = texts + + if len(generate_prompts) == 1: + prompt_kwargs = { + "text": generate_prompts[0], + "image_data": images[0], + } + else: + prompt_kwargs = { + "text": generate_prompts, + "image_data": images, + } + else: + # List of integers (token IDs) or empty list + prompt_kwargs = {"input_ids": prompt} + else: + # Other types (should not happen but handle gracefully) + prompt_kwargs = {"input_ids": prompt} + + adapted_request = EmbeddingReqInput( + **prompt_kwargs, + rid=request.rid, + ) + + return adapted_request, request + + async def _handle_non_streaming_request( + self, + adapted_request: EmbeddingReqInput, + request: EmbeddingRequest, + raw_request: Request, + ) -> Union[EmbeddingResponse, ErrorResponse, ORJSONResponse]: + """Handle the embedding request""" + try: + ret = await self.tokenizer_manager.generate_request( + adapted_request, raw_request + ).__anext__() + except ValueError as e: + return self.create_error_response(str(e)) + + if not isinstance(ret, list): + ret = [ret] + + response = self._build_embedding_response(ret) + return response + + def _build_embedding_response(self, ret: List[Dict[str, Any]]) -> EmbeddingResponse: + """Build the embedding response""" + embedding_objects = [] + prompt_tokens = 0 + + for idx, ret_item in enumerate(ret): + embedding_objects.append( + EmbeddingObject( + embedding=ret_item["embedding"], + index=idx, + ) + ) + # Handle missing prompt_tokens gracefully + meta_info = ret_item.get("meta_info", {}) + prompt_tokens += meta_info.get("prompt_tokens", 0) + + return EmbeddingResponse( + data=embedding_objects, + model=self.tokenizer_manager.model_path, + usage=UsageInfo( + prompt_tokens=prompt_tokens, + total_tokens=prompt_tokens, + ), + ) diff --git a/python/sglang/srt/entrypoints/openai/serving_rerank.py b/python/sglang/srt/entrypoints/openai/serving_rerank.py new file mode 100644 index 000000000..b053c55b3 --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/serving_rerank.py @@ -0,0 +1,102 @@ +import logging +from typing import Any, Dict, List, Optional, Union + +from fastapi import Request +from fastapi.responses import ORJSONResponse + +from sglang.srt.entrypoints.openai.protocol import ( + ErrorResponse, + RerankResponse, + V1RerankReqInput, +) +from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase +from sglang.srt.managers.io_struct import EmbeddingReqInput + +logger = logging.getLogger(__name__) + + +class OpenAIServingRerank(OpenAIServingBase): + """Handler for /v1/rerank requests""" + + # NOTE: /v1/rerank is not an official OpenAI endpoint. This module may be moved + # to another module in the future. + + def _request_id_prefix(self) -> str: + return "rerank-" + + def _validate_request(self, request: V1RerankReqInput) -> Optional[str]: + """Validate rerank request format and content""" + if not request.query: + return "Query cannot be empty" + + if isinstance(request.query, str): + if not request.query.strip(): + return "Query cannot be empty or whitespace only" + + if not request.documents: + return "Documents cannot be empty" + + for doc in request.documents: + if not doc: + return "Each document must be a non-empty string" + if isinstance(doc, str) and not doc.strip(): + return "Each document cannot be empty or whitespace only" + + return None + + def _convert_to_internal_request( + self, request: V1RerankReqInput + ) -> tuple[EmbeddingReqInput, V1RerankReqInput]: + """Convert OpenAI rerank request to internal embedding format""" + # Create pairs of [query, document] for each document + pairs = [] + for doc in request.documents: + pairs.append([request.query, doc]) + + adapted_request = EmbeddingReqInput( + text=pairs, + is_cross_encoder_request=True, + ) + + return adapted_request, request + + async def _handle_non_streaming_request( + self, + adapted_request: EmbeddingReqInput, + request: V1RerankReqInput, + raw_request: Request, + ) -> Union[List[RerankResponse], ErrorResponse, ORJSONResponse]: + """Handle the rerank request""" + try: + ret = await self.tokenizer_manager.generate_request( + adapted_request, raw_request + ).__anext__() + + except ValueError as e: + return self.create_error_response(str(e)) + + if not isinstance(ret, list): + ret = [ret] + + responses = self._build_rerank_response(ret, request) + return responses + + def _build_rerank_response( + self, ret: List[Dict[str, Any]], request: V1RerankReqInput + ) -> List[RerankResponse]: + """Build the rerank response from generation results""" + responses = [] + for idx, ret_item in enumerate(ret): + responses.append( + RerankResponse( + score=ret_item["embedding"], + document=request.documents[idx], + index=idx, + meta_info=ret_item["meta_info"], + ) + ) + + # Sort by score in descending order (highest relevance first) + responses.sort(key=lambda x: x.score, reverse=True) + + return responses diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py new file mode 100644 index 000000000..3f7619678 --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -0,0 +1,1276 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vLLM's OpenAIServingResponses +"""Handler for /v1/responses requests""" +from __future__ import annotations + +import asyncio +import copy +import json +import logging +import time +from contextlib import AsyncExitStack +from http import HTTPStatus +from typing import TYPE_CHECKING, Any, AsyncGenerator, AsyncIterator, Optional, Union + +import jinja2 +import openai.types.responses as openai_responses_types +from fastapi import Request +from fastapi.responses import ORJSONResponse +from openai.types.responses import ( + ResponseOutputMessage, + ResponseOutputText, + ResponseReasoningItem, +) +from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall +from openai.types.responses.response_reasoning_item import ( + Content as ResponseReasoningTextContent, +) +from openai_harmony import Message as OpenAIMessage + +from sglang.srt.entrypoints.context import ( + ConversationContext, + HarmonyContext, + SimpleContext, + StreamingHarmonyContext, +) +from sglang.srt.entrypoints.harmony_utils import ( + get_developer_message, + get_stop_tokens_for_assistant_actions, + get_system_message, + get_user_message, + parse_output_message, + parse_remaining_state, + parse_response_input, + render_for_completion, +) +from sglang.srt.entrypoints.openai.protocol import ( + ChatCompletionMessageParam, + ChatCompletionRequest, + PromptTokenUsageInfo, + RequestResponseMetadata, + ResponsesRequest, + ResponsesResponse, + UsageInfo, +) +from sglang.srt.entrypoints.openai.serving_chat import OpenAIServingChat +from sglang.srt.entrypoints.openai.tool_server import MCPToolServer, ToolServer +from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.parser.reasoning_parser import ReasoningParser +from sglang.srt.utils import random_uuid + +if TYPE_CHECKING: + from sglang.srt.managers.template_manager import TemplateManager + from sglang.srt.managers.tokenizer_manager import TokenizerManager + +logger = logging.getLogger(__name__) + + +class OpenAIServingResponses(OpenAIServingChat): + """Handler for /v1/responses requests""" + + def __init__( + self, + tokenizer_manager: TokenizerManager, + template_manager: TemplateManager, + *, + enable_prompt_tokens_details: bool = False, + enable_force_include_usage: bool = False, + tool_server: Optional[ToolServer] = None, + ) -> None: + super().__init__(tokenizer_manager, template_manager) + + # template_manager is already set by parent class + self.reasoning_parser = self.tokenizer_manager.server_args.reasoning_parser + self.enable_prompt_tokens_details = enable_prompt_tokens_details + self.enable_force_include_usage = enable_force_include_usage + + # Get default sampling params from model config if available + self.default_sampling_params = {} + + self.supports_browsing = ( + tool_server.has_tool("browser") if tool_server else False + ) + self.supports_code_interpreter = ( + tool_server.has_tool("python") if tool_server else False + ) + self.tool_server = tool_server + # Get from model config + self.use_harmony = ( + self.tokenizer_manager.model_config.hf_config.model_type == "gpt_oss" + ) + + if self.use_harmony: + # OpenAI models have two EOS-like tokens: <|return|> and <|call|>. + # We need to add them to the stop token ids. + if "stop_token_ids" not in self.default_sampling_params: + self.default_sampling_params["stop_token_ids"] = [] + self.default_sampling_params["stop_token_ids"].extend( + get_stop_tokens_for_assistant_actions() + ) + + # Response storage for background and retrieval operations + # Note: In production, this should use a proper storage backend (Redis, database) + # with TTL/expiration to prevent memory leaks + self.response_store: dict[str, ResponsesResponse] = {} + self.response_store_lock = asyncio.Lock() + + # Message storage for conversation continuity + # Note: In production, this should use a proper storage backend (Redis, database) + # with TTL/expiration to prevent memory leaks + self.msg_store: dict[ + str, Union[list[ChatCompletionMessageParam], list["OpenAIMessage"]] + ] = {} + + self.background_tasks: dict[str, asyncio.Task] = {} + + def _request_id_prefix(self) -> str: + return "resp_" + + async def create_responses( + self, + request: ResponsesRequest, + raw_request: Optional[Request] = None, + ) -> Union[AsyncGenerator[str, None], ResponsesResponse, ORJSONResponse]: + # Validate model + if not self.tokenizer_manager: + return self.create_error_response("Model not loaded") + + # FIXME: If the engine is dead, raise an error + # This is required for the streaming case + + # Handle the previous response ID + prev_response_id = request.previous_response_id + if prev_response_id is not None: + if not prev_response_id.startswith("resp_"): + return self._make_invalid_id_error(prev_response_id) + async with self.response_store_lock: + prev_response = self.response_store.get(prev_response_id) + if prev_response is None: + return self._make_not_found_error(prev_response_id) + else: + prev_response = None + + try: + model_name = request.model + tokenizer = self.tokenizer_manager.tokenizer + + if self.use_harmony: + messages, request_prompts, engine_prompts = ( + self._make_request_with_harmony(request, prev_response) + ) + else: + messages, request_prompts, engine_prompts = await self._make_request( + request, prev_response, tokenizer + ) + + except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e: + logger.exception("Error in preprocessing prompt inputs") + return self.create_error_response(f"{e} {e.__cause__}") + + request_metadata = RequestResponseMetadata(request_id=request.request_id) + if raw_request: + raw_request.state.request_metadata = request_metadata + + if ( + self.tool_server is not None + and isinstance(self.tool_server, MCPToolServer) + and (request.background or request.stream) + and request.tools + and any( + tool.type in ["web_search_preview", "code_interpreter"] + for tool in request.tools + ) + ): + return self.create_error_response( + "MCP tool server is not supported in background mode and " + "streaming mode" + ) + + # Schedule the request and get the result generator + generators: list[AsyncGenerator[Any, None]] = [] + tool_list = [] + if self.use_harmony: + if self.supports_browsing: + tool_list.append("browser") + if self.supports_code_interpreter: + tool_list.append("python") + async with AsyncExitStack() as exit_stack: + try: + if self.tool_server is not None: + tool_session_ctxs: dict[str, Any] = { + tool_name: exit_stack.enter_async_context( + self.tool_server.get_tool_session(tool_name) + ) + for tool_name in tool_list + } + tool_sessions = {} + for tool_name in tool_list: + tool_sessions[tool_name] = await tool_session_ctxs[tool_name] + else: + assert len(tool_list) == 0 + tool_sessions = {} + for i, engine_prompt in enumerate(engine_prompts): + # Calculate default max tokens from context length minus prompt length + if hasattr(engine_prompt, "__len__"): + prompt_length = len(engine_prompt) + elif isinstance(engine_prompt, list): + prompt_length = len(engine_prompt) + else: + prompt_length = 0 + + context_len = ( + self.tokenizer_manager.model_config.context_len + if hasattr(self.tokenizer_manager.model_config, "context_len") + else 4096 + ) + default_max_tokens = max( + context_len - prompt_length, 512 + ) # Ensure minimum 512 tokens + sampling_params = request.to_sampling_params( + default_max_tokens, self.default_sampling_params + ) + + context: ConversationContext + if self.use_harmony: + if request.stream: + context = StreamingHarmonyContext(messages, tool_sessions) + else: + context = HarmonyContext(messages, tool_sessions) + else: + context = SimpleContext() + + # Create GenerateReqInput for SGLang + adapted_request = GenerateReqInput( + input_ids=engine_prompt, + sampling_params=sampling_params, + stream=request.stream, + rid=request.request_id, + background=request.background, + ) + + generator = self._generate_with_builtin_tools( + request.request_id, + request_prompts[i], + adapted_request, + sampling_params, + context, + raw_request=raw_request, + priority=request.priority, + ) + generators.append(generator) + except ValueError as e: + return self.create_error_response(str(e)) + + assert len(generators) == 1 + (result_generator,) = generators + + # Store the input messages + if request.store: + self.msg_store[request.request_id] = messages + + if request.background: + created_time = int(time.time()) + response = ResponsesResponse.from_request( + request, + sampling_params, + model_name=model_name, + created_time=created_time, + output=[], + status="queued", + usage=None, + ) + async with self.response_store_lock: + self.response_store[response.id] = response + + # Run the request in the background + task = asyncio.create_task( + self._run_background_request( + request, + sampling_params, + result_generator, + context, + model_name, + tokenizer, + request_metadata, + created_time, + ), + name=f"create_{response.id}", + ) + + # For cleanup + self.background_tasks[response.id] = task + task.add_done_callback( + lambda _: self.background_tasks.pop(response.id, None) + ) + return response + + if request.stream: + return self.responses_stream_generator( + request, + sampling_params, + result_generator, + context, + model_name, + tokenizer, + request_metadata, + ) + try: + result: Union[ORJSONResponse, ResponsesResponse] = ( + await self.responses_full_generator( + request, + sampling_params, + result_generator, + context, + model_name, + tokenizer, + request_metadata, + ) + ) + return result + except Exception as e: + return self.create_error_response(str(e)) + return self.create_error_response("Unknown error") + + async def _make_request( + self, + request: ResponsesRequest, + prev_response: Optional[ResponsesResponse], + tokenizer: Any, + ): + # Construct the input messages + messages = self._construct_input_messages(request, prev_response) + + # Follow SGLang's pattern: create a ChatCompletionRequest and process messages + try: + # Convert ResponsesRequest to ChatCompletionRequest for processing + chat_request = ChatCompletionRequest( + model=request.model, + messages=messages, + stream=request.stream, + ) + + # Follow SGLang's _process_messages pattern + is_multimodal = self.tokenizer_manager.model_config.is_multimodal + processed_messages = self._process_messages(chat_request, is_multimodal) + + # Extract the results + if is_multimodal: + request_prompts = [processed_messages.prompt] + engine_prompts = [processed_messages.prompt] + else: + request_prompts = [processed_messages.prompt_ids] + engine_prompts = [processed_messages.prompt_ids] + + except Exception as e: + logger.warning(f"Chat processing failed, using fallback: {e}") + # Fallback to simple encoding + prompt_text = "" + for msg in messages: + role = msg.get("role", "user") + content = msg.get("content", "") + prompt_text += f"{role}: {content}\n" + prompt_ids = tokenizer.encode(prompt_text) + request_prompts = [prompt_ids] + engine_prompts = [prompt_ids] + + return messages, request_prompts, engine_prompts + + def _make_request_with_harmony( + self, + request: ResponsesRequest, + prev_response: Optional[ResponsesResponse], + ): + if request.tool_choice != "auto": + raise NotImplementedError( + "Only 'auto' tool_choice is supported in " "response API" + ) + messages = self._construct_input_messages_with_harmony(request, prev_response) + prompt_token_ids = render_for_completion(messages) + engine_prompt = prompt_token_ids + return messages, [prompt_token_ids], [engine_prompt] + + async def responses_full_generator( + self, + request: ResponsesRequest, + sampling_params: Any, + result_generator: AsyncIterator[Any], + context: ConversationContext, + model_name: str, + tokenizer: Any, + request_metadata: RequestResponseMetadata, + created_time: Optional[int] = None, + ) -> Union[ResponsesResponse, ORJSONResponse]: + if created_time is None: + created_time = int(time.time()) + + try: + async for _ in result_generator: + pass + except asyncio.CancelledError: + return self.create_error_response("Client disconnected") + except ValueError as e: + return self.create_error_response(str(e)) + + if self.use_harmony: + assert isinstance(context, HarmonyContext) + output = self._make_response_output_items_with_harmony(context) + # TODO: these are all 0 for now! + num_prompt_tokens = context.num_prompt_tokens + num_generated_tokens = context.num_output_tokens + num_cached_tokens = context.num_cached_tokens + num_reasoning_tokens = context.num_reasoning_tokens + else: + assert isinstance(context, SimpleContext) + final_res = context.last_output + assert final_res is not None + + output = self._make_response_output_items( + request, final_res["text"], tokenizer + ) + + # Calculate usage from actual output + if hasattr(final_res, "meta_info"): + num_prompt_tokens = final_res.meta_info.get("prompt_tokens", 0) + num_generated_tokens = final_res.meta_info.get("completion_tokens", 0) + num_cached_tokens = final_res.meta_info.get("cached_tokens", 0) + elif hasattr(final_res, "prompt_token_ids") and hasattr( + final_res, "outputs" + ): + # Fallback calculation if meta_info not available + num_prompt_tokens = ( + len(final_res.prompt_token_ids) if final_res.prompt_token_ids else 0 + ) + num_generated_tokens = ( + len(final_res.outputs[0].token_ids) + if final_res.outputs and final_res.outputs[0].token_ids + else 0 + ) + num_cached_tokens = getattr(final_res, "num_cached_tokens", 0) + num_reasoning_tokens = 0 + else: + # Final fallback + num_prompt_tokens = 0 + num_generated_tokens = 0 + num_cached_tokens = 0 + num_reasoning_tokens = 0 + + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=num_generated_tokens, + total_tokens=num_prompt_tokens + num_generated_tokens, + reasoning_tokens=num_reasoning_tokens, + ) + if self.enable_prompt_tokens_details and num_cached_tokens: + usage.prompt_tokens_details = PromptTokenUsageInfo( + cached_tokens=num_cached_tokens + ) + request_metadata.final_usage_info = usage + + response = ResponsesResponse.from_request( + request, + sampling_params, + model_name=model_name, + created_time=created_time, + output=output, + status="completed", + usage=usage, + ) + + if request.store: + async with self.response_store_lock: + stored_response = self.response_store.get(response.id) + # If the response is already cancelled, don't update it + if stored_response is None or stored_response.status != "cancelled": + self.response_store[response.id] = response + + return response + + def _make_response_output_items( + self, + request: ResponsesRequest, + final_output: Any, + tokenizer: Any, + ): + # Handle reasoning parsing if enabled + if self.reasoning_parser: + # Use standard reasoning parser (openai maps to T4Detector internally) + reasoning_parser = ReasoningParser( + model_type=self.reasoning_parser, stream_reasoning=False + ) + reasoning_content, content = reasoning_parser.parse_non_stream(final_output) + else: + reasoning_content = None + content = final_output + + output_items = [] + if reasoning_content: + reasoning_item = ResponseReasoningItem( + id=f"rs_{random_uuid()}", + type="reasoning", + summary=[], + content=[ + ResponseReasoningTextContent( + type="reasoning_text", text=reasoning_content + ), + ], + status=None, + ) + output_items.append(reasoning_item) + if content: + output_text = ResponseOutputText( + text=content, + annotations=[], # TODO + type="output_text", + logprobs=None, # TODO + ) + message = ResponseOutputMessage( + id=f"msg_{random_uuid()}", + content=[output_text], + role="assistant", + status="completed", + type="message", + ) + output_items.append(message) + return output_items + + def _make_response_output_items_with_harmony( + self, + context: HarmonyContext, + ): + output_items = [] + num_init_messages = context.num_init_messages + for msg in context.messages[num_init_messages:]: + output_items.extend(parse_output_message(msg)) + # Handle the generation stopped in the middle (if any). + last_items = parse_remaining_state(context.parser) + if last_items: + output_items.extend(last_items) + return output_items + + def _construct_input_messages( + self, + request: ResponsesRequest, + prev_response: Optional[ResponsesResponse] = None, + ) -> list[ChatCompletionMessageParam]: + messages: list[ChatCompletionMessageParam] = [] + if request.instructions: + messages.append( + { + "role": "system", + "content": request.instructions, + } + ) + + # Prepend the conversation history + if prev_response is not None: + # Add the previous messages + prev_msg = self.msg_store[prev_response.id] + messages.extend(prev_msg) + + # Add the previous output + for output_item in prev_response.output: + # NOTE: We skip the reasoning output of the previous response + if isinstance(output_item, ResponseReasoningItem): + continue + for content in output_item.content: + messages.append( + { + "role": "system", + "content": request.instructions, + } + ) + + # Append the new input + # Responses API supports simple text inputs without chat format + if isinstance(request.input, str): + messages.append({"role": "user", "content": request.input}) + else: + messages.extend(request.input) # type: ignore + return messages + + def _construct_input_messages_with_harmony( + self, + request: ResponsesRequest, + prev_response: Optional[ResponsesResponse], + ) -> list["OpenAIMessage"]: + messages: list["OpenAIMessage"] = [] + if prev_response is None: + # New conversation. + reasoning_effort = request.reasoning.effort if request.reasoning else None + tool_types = [tool.type for tool in request.tools] + enable_browser = ( + "web_search_preview" in tool_types and self.tool_server is not None + ) + enable_code_interpreter = ( + "code_interpreter" in tool_types and self.tool_server is not None + ) + sys_msg = get_system_message( + reasoning_effort=reasoning_effort, + browser_description=( + self.tool_server.get_tool_description("browser") + if self.tool_server and enable_browser + else None + ), + python_description=( + self.tool_server.get_tool_description("python") + if self.tool_server and enable_code_interpreter + else None + ), + ) + messages.append(sys_msg) + dev_msg = get_developer_message(request.instructions, request.tools) + messages.append(dev_msg) + else: + # Continue the previous conversation. + # FIXME: Currently, request params like reasoning and + # instructions are ignored. + prev_msgs = self.msg_store[prev_response.id] + # Remove the previous chain-of-thoughts if there is a new "final" + # message. + if ( + len(prev_msgs) > 0 + and hasattr(prev_msgs[-1], "channel") + and prev_msgs[-1].channel == "final" + ): # type: ignore[union-attr] + prev_final_msg_idx = -1 + for i in range(len(prev_msgs) - 2, -1, -1): + if ( + hasattr(prev_msgs[i], "channel") + and prev_msgs[i].channel == "final" + ): # type: ignore[union-attr] + prev_final_msg_idx = i + break + recent_turn_msgs = prev_msgs[prev_final_msg_idx + 1 :] + del prev_msgs[prev_final_msg_idx + 1 :] + for msg in recent_turn_msgs: + if ( + hasattr(msg, "channel") and msg.channel != "analysis" + ): # type: ignore[union-attr] + prev_msgs.append(msg) + messages.extend(prev_msgs) + # Append the new input. + # Responses API supports simple text inputs without chat format. + if isinstance(request.input, str): + messages.append(get_user_message(request.input)) + else: + if prev_response is not None: + prev_outputs = copy(prev_response.output) + else: + prev_outputs = [] + for response_msg in request.input: + messages.append(parse_response_input(response_msg, prev_outputs)) + if isinstance(response_msg, ResponseFunctionToolCall): + prev_outputs.append(response_msg) + return messages + + async def _run_background_request( + self, + request: ResponsesRequest, + sampling_params: Any, + result_generator: AsyncIterator[Any], + context: ConversationContext, + model_name: str, + tokenizer: Any, + request_metadata: RequestResponseMetadata, + created_time: Optional[int] = None, + *args, + **kwargs, + ): + try: + # Update the status to "in_progress" + async with self.response_store_lock: + stored_response = self.response_store.get(request.request_id) + assert stored_response is not None + stored_response.status = "in_progress" + + response = await self.responses_full_generator( + request, + sampling_params, + result_generator, + context, + model_name, + tokenizer, + request_metadata, + created_time, + *args, + **kwargs, + ) + except Exception as e: + logger.exception("Background request failed for %s", request.request_id) + response = self.create_error_response(str(e)) + + if isinstance(response, ORJSONResponse): + # If the request has failed, update the status to "failed" + response_id = request.request_id + async with self.response_store_lock: + stored_response = self.response_store.get(response_id) + assert stored_response is not None + if stored_response.status not in ("completed", "cancelled"): + stored_response.status = "failed" + + async def retrieve_responses( + self, + response_id: str, + ) -> Union[ResponsesResponse, ORJSONResponse]: + if not response_id.startswith("resp_"): + return self._make_invalid_id_error(response_id) + + async with self.response_store_lock: + response = self.response_store.get(response_id) + + if response is None: + return self._make_not_found_error(response_id) + return response + + async def cancel_responses( + self, + response_id: str, + ) -> Union[ResponsesResponse, ORJSONResponse]: + if not response_id.startswith("resp_"): + return self._make_invalid_id_error(response_id) + + async with self.response_store_lock: + response = self.response_store.get(response_id) + if response is None: + return self._make_not_found_error(response_id) + + prev_status = response.status + if prev_status not in ("queued", "in_progress"): + return self.create_error_response( + err_type="invalid_request_error", + message="Cannot cancel a synchronous response.", + ) + + # Update the status to "cancelled" + response.status = "cancelled" + + # Abort the request + if task := self.background_tasks.get(response_id): + task.cancel() + try: + await task + except asyncio.CancelledError: + logger.exception("Background task for %s was cancelled", response_id) + return response + + def _make_invalid_id_error(self, response_id: str): + return self.create_error_response( + message=( + f"Invalid 'response_id': '{response_id}'. " + "Expected an ID that begins with 'resp'." + ), + err_type="invalid_request_error", + param="response_id", + ) + + def _make_not_found_error(self, response_id: str): + return self.create_error_response( + message=f"Response with id '{response_id}' not found.", + err_type="invalid_request_error", + status_code=HTTPStatus.NOT_FOUND, + param="response_id", + ) + + async def responses_stream_generator( + self, + request: ResponsesRequest, + sampling_params: Any, + result_generator: AsyncIterator[StreamingHarmonyContext], + context: StreamingHarmonyContext, + model_name: str, + tokenizer: Any, + request_metadata: RequestResponseMetadata, + created_time: Optional[int] = None, + ) -> AsyncGenerator[str, None]: + # TODO: + # 1. Handle disconnect + + created_time = created_time or int(time.time()) + + sequence_number = 0 + + def _send_event(event): + nonlocal sequence_number + # Set sequence_number if the event has this attribute + if hasattr(event, "sequence_number"): + event.sequence_number = sequence_number + sequence_number += 1 + # Get event type from the event's type field if it exists + event_type = getattr(event, "type", "unknown") + return ( + f"event: {event_type}\n" + f"data: {event.model_dump_json(indent=None)}\n\n" + ) + + current_content_index = 0 + current_output_index = 0 + current_item_id = f"item_{random_uuid()}" + sent_output_item_added = False + + initial_response = ResponsesResponse.from_request( + request, + sampling_params, + model_name=model_name, + created_time=created_time, + output=[], + status="in_progress", + usage=None, + ).model_dump() + yield _send_event( + openai_responses_types.ResponseCreatedEvent( + type="response.created", + sequence_number=-1, + response=initial_response, + ) + ) + yield _send_event( + openai_responses_types.ResponseInProgressEvent( + type="response.in_progress", + sequence_number=-1, + response=initial_response, + ) + ) + + async for ctx in result_generator: + + if ctx.is_expecting_start(): + current_output_index += 1 + sent_output_item_added = False + + if len(ctx.parser.messages) > 0: + previous_item = ctx.parser.messages[-1] + if previous_item.recipient is not None: + # Deal with tool call here + pass + elif previous_item.channel == "analysis": + reasoning_item = ResponseReasoningItem( + id=f"rs_{random_uuid()}", + type="reasoning", + summary=[], + content=[ + ResponseReasoningTextContent( + text=previous_item.content[0].text, + type="reasoning_text", + ), + ], + status="completed", + ) + yield _send_event( + openai_responses_types.ResponseReasoningTextDoneEvent( + type="response.reasoning_text.done", + item_id=current_item_id, + sequence_number=-1, + output_index=current_output_index, + content_index=current_content_index, + text=previous_item.content[0].text, + ) + ) + yield _send_event( + openai_responses_types.ResponseOutputItemDoneEvent( + type="response.output_item.done", + sequence_number=-1, + output_index=current_output_index, + item=reasoning_item, + ) + ) + elif previous_item.channel == "final": + text_content = openai_responses_types.ResponseOutputText( + type="output_text", + text=previous_item.content[0].text, + annotations=[], + ) + yield _send_event( + openai_responses_types.ResponseTextDoneEvent( + type="response.output_text.done", + sequence_number=-1, + output_index=current_output_index, + content_index=current_content_index, + text=previous_item.content[0].text, + logprobs=[], + item_id=current_item_id, + ) + ) + yield _send_event( + openai_responses_types.ResponseContentPartDoneEvent( + type="response.content_part.done", + sequence_number=-1, + item_id=current_item_id, + output_index=current_output_index, + content_index=current_content_index, + part=text_content, + ) + ) + yield _send_event( + openai_responses_types.ResponseOutputItemDoneEvent( + type="response.output_item.done", + sequence_number=-1, + output_index=current_output_index, + item=openai_responses_types.ResponseOutputMessage( + id=current_item_id, + type="message", + role="assistant", + content=[text_content], + status="completed", + ), + ) + ) + + if ctx.parser.last_content_delta: + if ( + ctx.parser.current_channel == "final" + and ctx.parser.current_recipient is None + ): + if not sent_output_item_added: + sent_output_item_added = True + yield _send_event( + openai_responses_types.ResponseOutputItemAddedEvent( + type="response.output_item.added", + sequence_number=-1, + output_index=current_output_index, + item=openai_responses_types.ResponseOutputMessage( + id=current_item_id, + type="message", + role="assistant", + content=[], + status="in_progress", + ), + ) + ) + yield _send_event( + openai_responses_types.ResponseContentPartAddedEvent( + type="response.content_part.added", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + content_index=current_content_index, + part=openai_responses_types.ResponseOutputText( + type="output_text", + text="", + annotations=[], + logprobs=None, + ), + ) + ) + yield _send_event( + openai_responses_types.ResponseTextDeltaEvent( + type="response.output_text.delta", + sequence_number=-1, + content_index=current_content_index, + output_index=current_output_index, + item_id=current_item_id, + delta=ctx.parser.last_content_delta, + # TODO, use logprobs from ctx.last_request_output + logprobs=[], + ) + ) + elif ( + ctx.parser.current_channel == "analysis" + and ctx.parser.current_recipient is None + ): + if not sent_output_item_added: + sent_output_item_added = True + yield _send_event( + openai_responses_types.ResponseOutputItemAddedEvent( + type="response.output_item.added", + sequence_number=-1, + output_index=current_output_index, + item=openai_responses_types.ResponseReasoningItem( + type="reasoning", + id=current_item_id, + summary=[], + status="in_progress", + ), + ) + ) + yield _send_event( + openai_responses_types.ResponseContentPartAddedEvent( + type="response.content_part.added", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + content_index=current_content_index, + # TODO: migrate this to + # ResponseReasoningTextContent for now + part=openai_responses_types.ResponseOutputText( + type="output_text", + text="", + annotations=[], + logprobs=None, + ), + ) + ) + # TODO: migrate to OpenAI types once updated. + yield _send_event( + openai_responses_types.ResponseReasoningTextDeltaEvent( + type="response.reasoning_text.delta", + item_id=current_item_id, + output_index=current_output_index, + content_index=current_content_index, + delta=ctx.parser.last_content_delta, + sequence_number=-1, + ) + ) + + if ctx.is_assistant_action_turn() and len(ctx.parser.messages) > 0: + previous_item = ctx.parser.messages[-1] + if ( + self.supports_browsing + and previous_item.recipient is not None + and previous_item.recipient.startswith("browser.") + ): + function_name = previous_item.recipient[len("browser.") :] + action = None + parsed_args = json.loads(previous_item.content[0].text) + if function_name == "search": + action = openai_responses_types.response_function_web_search.ActionSearch( + type="search", + query=parsed_args["query"], + ) + elif function_name == "open": + action = openai_responses_types.response_function_web_search.ActionOpenPage( + type="open_page", + # TODO: translate to url + url=f"cursor:{parsed_args.get('cursor', '')}", + ) + elif function_name == "find": + action = openai_responses_types.response_function_web_search.ActionFind( + type="find", + pattern=parsed_args["pattern"], + # TODO: translate to url + url=f"cursor:{parsed_args.get('cursor', '')}", + ) + else: + raise ValueError(f"Unknown function name: {function_name}") + + yield _send_event( + openai_responses_types.ResponseOutputItemAddedEvent( + type="response.output_item.added", + sequence_number=-1, + output_index=current_output_index, + item=openai_responses_types.response_function_web_search.ResponseFunctionWebSearch( + # TODO: generate a unique id for web search call + type="web_search_call", + id=current_item_id, + action=action, + status="in_progress", + ), + ) + ) + yield _send_event( + openai_responses_types.ResponseWebSearchCallInProgressEvent( + type="response.web_search_call.in_progress", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + ) + ) + yield _send_event( + openai_responses_types.ResponseWebSearchCallSearchingEvent( + type="response.web_search_call.searching", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + ) + ) + + # enqueue + yield _send_event( + openai_responses_types.ResponseWebSearchCallCompletedEvent( + type="response.web_search_call.completed", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + ) + ) + yield _send_event( + openai_responses_types.ResponseOutputItemDoneEvent( + type="response.output_item.done", + sequence_number=-1, + output_index=current_output_index, + item=openai_responses_types.ResponseFunctionWebSearch( + type="web_search_call", + id=current_item_id, + action=action, + status="completed", + ), + ) + ) + + if ( + self.supports_code_interpreter + and previous_item.recipient is not None + and previous_item.recipient.startswith("python") + ): + yield _send_event( + openai_responses_types.ResponseOutputItemAddedEvent( + type="response.output_item.added", + sequence_number=-1, + output_index=current_output_index, + item=openai_responses_types.ResponseCodeInterpreterToolCallParam( + type="code_interpreter_call", + id=current_item_id, + code="", + container_id="auto", + outputs=[], + status="in_progress", + ), + ) + ) + yield _send_event( + openai_responses_types.ResponseCodeInterpreterCallInProgressEvent( + type="response.code_interpreter_call.in_progress", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + ) + ) + # TODO: do we need to add delta event here? + yield _send_event( + openai_responses_types.ResponseCodeInterpreterCallCodeDoneEvent( + type="response.code_interpreter_call_code.done", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + code=previous_item.content[0].text, + ) + ) + yield _send_event( + openai_responses_types.ResponseCodeInterpreterCallInterpretingEvent( + type="response.code_interpreter_call.interpreting", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + ) + ) + yield _send_event( + openai_responses_types.ResponseCodeInterpreterCallCompletedEvent( + type="response.code_interpreter_call.completed", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + ) + ) + yield _send_event( + openai_responses_types.ResponseOutputItemDoneEvent( + type="response.output_item.done", + sequence_number=-1, + output_index=current_output_index, + item=openai_responses_types.ResponseCodeInterpreterToolCallParam( + type="code_interpreter_call", + id=current_item_id, + code=previous_item.content[0].text, + container_id="auto", + # TODO: add outputs here + outputs=[], + status="completed", + ), + ) + ) + + async def empty_async_generator(): + if False: + yield + + final_response = await self.responses_full_generator( + request, + sampling_params, + empty_async_generator(), + context, + model_name, + tokenizer, + request_metadata, + created_time=created_time, + ) + # Convert final_response to the format expected by ResponseCompletedEvent + response_dict = final_response.model_dump() + + # Convert UsageInfo to ResponseUsage format + if response_dict.get("usage"): + usage_info = response_dict["usage"] + response_dict["usage"] = { + "input_tokens": usage_info.get("prompt_tokens", 0), + "input_tokens_details": { + "cached_tokens": usage_info.get("cached_tokens", 0) + }, + "output_tokens": usage_info.get("completion_tokens", 0), + "output_tokens_details": { + "reasoning_tokens": usage_info.get("reasoning_tokens", 0) + }, + "total_tokens": usage_info.get("total_tokens", 0), + } + + yield _send_event( + openai_responses_types.ResponseCompletedEvent( + type="response.completed", + sequence_number=-1, + response=response_dict, + ) + ) + + async def _generate_with_builtin_tools( + self, + request_id: str, + request_prompt: Any, + adapted_request: GenerateReqInput, + sampling_params: Any, + context: ConversationContext, + raw_request: Optional[Request] = None, + priority: Optional[int] = None, + **kwargs, + ) -> AsyncGenerator[Any, None]: + """Generate with builtin tool support for harmony-based models.""" + orig_priority = priority or 0 + + while True: + # Generate using SGLang's tokenizer manager + generator = self.tokenizer_manager.generate_request( + adapted_request, raw_request + ) + + async for res in generator: + context.append_output(res) + # NOTE(woosuk): The stop condition is handled by the engine. + yield context + + if not context.need_builtin_tool_call(): + # The model did not ask for a tool call, so we're done. + break + + # Call the tool and update the context with the result. + tool_output = await context.call_tool() + context.append_output(tool_output) + + # Prepare for the next generation turn + # Render the updated conversation for the next completion + prompt_token_ids = context.render_for_completion() + + # Update the adapted request with new prompt + adapted_request = GenerateReqInput( + input_ids=prompt_token_ids, + sampling_params=sampling_params, + stream=adapted_request.stream, + rid=request_id, + return_logprob=adapted_request.return_logprob, + logprob_start_len=adapted_request.logprob_start_len, + top_logprobs_num=adapted_request.top_logprobs_num, + return_text_in_logprobs=adapted_request.return_text_in_logprobs, + return_hidden_states=adapted_request.return_hidden_states, + background=adapted_request.background, + ) + + # Update sampling params with reduced max_tokens + if hasattr(sampling_params, "max_new_tokens") or isinstance( + sampling_params, dict + ): + context_len = getattr( + self.tokenizer_manager.model_config, "context_len", 4096 + ) + remaining_tokens = context_len - len(prompt_token_ids) - 1 + + if isinstance(sampling_params, dict): + sampling_params["max_new_tokens"] = max(remaining_tokens, 1) + else: + sampling_params.max_new_tokens = max(remaining_tokens, 1) + + # Slightly reduce priority for subsequent tool calls + priority = orig_priority - 1 diff --git a/python/sglang/srt/entrypoints/openai/serving_score.py b/python/sglang/srt/entrypoints/openai/serving_score.py new file mode 100644 index 000000000..fc8ce5dca --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/serving_score.py @@ -0,0 +1,61 @@ +import logging +from typing import Union + +from fastapi import Request + +from sglang.srt.entrypoints.openai.protocol import ( + ErrorResponse, + ScoringRequest, + ScoringResponse, +) +from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase + +logger = logging.getLogger(__name__) + + +class OpenAIServingScore(OpenAIServingBase): + """Handler for /v1/score requests""" + + # NOTE: /v1/rerank is not an official OpenAI endpoint. This module may be moved + # to another module in the future. + + def _request_id_prefix(self) -> str: + return "score-" + + def _convert_to_internal_request( + self, + request: ScoringRequest, + ) -> tuple[ScoringRequest, ScoringRequest]: + """Convert OpenAI scoring request to internal format""" + # For scoring, we pass the request directly as the tokenizer_manager + # has a specialized score_request method that doesn't use GenerateReqInput + + return request, request + + async def _handle_non_streaming_request( + self, + adapted_request: ScoringRequest, + request: ScoringRequest, + raw_request: Request, + ) -> Union[ScoringResponse, ErrorResponse]: + """Handle the scoring request""" + try: + # Use tokenizer_manager's score_request method directly + scores = await self.tokenizer_manager.score_request( + query=request.query, + items=request.items, + label_token_ids=request.label_token_ids, + apply_softmax=request.apply_softmax, + item_first=request.item_first, + request=raw_request, + ) + + # Create response with just the scores, without usage info + response = ScoringResponse( + scores=scores, + model=request.model, + ) + return response + + except ValueError as e: + return self.create_error_response(str(e)) diff --git a/python/sglang/srt/entrypoints/openai/tool_server.py b/python/sglang/srt/entrypoints/openai/tool_server.py new file mode 100644 index 000000000..269d9e99e --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/tool_server.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging +from abc import ABC, abstractmethod +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from typing import Any + +try: + from mcp import ClientSession + from mcp.client.sse import sse_client + from mcp.types import ListToolsResult +except ImportError as e: + ClientSession = sse_client = ListToolsResult = e + +from openai_harmony import ToolDescription, ToolNamespaceConfig + +logger = logging.getLogger(__name__) + + +async def list_server_and_tools(server_url: str): + + async with sse_client(url=server_url) as streams, ClientSession( + *streams + ) as session: + initialize_response = await session.initialize() + list_tools_response = await session.list_tools() + return initialize_response, list_tools_response + + +def trim_schema(schema: dict) -> dict: + # Turn JSON Schema from MCP generated into Harmony's variant. + if "title" in schema: + del schema["title"] + if "default" in schema and schema["default"] is None: + del schema["default"] + if "anyOf" in schema: + # Turn "anyOf": [{"type": "type-1"}, {"type": "type-2"}] + # into "type": ["type-1", "type-2"] + # if there's more than 1 types, also remove "null" type as Harmony will + # just ignore it + types = [ + type_dict["type"] + for type_dict in schema["anyOf"] + if type_dict["type"] != "null" + ] + schema["type"] = types + del schema["anyOf"] + if "properties" in schema: + schema["properties"] = { + k: trim_schema(v) for k, v in schema["properties"].items() + } + return schema + + +def post_process_tools_description( + list_tools_result: "ListToolsResult", +) -> "ListToolsResult": + # Adapt the MCP tool result for Harmony + for tool in list_tools_result.tools: + tool.inputSchema = trim_schema(tool.inputSchema) + + # Some tools schema don't need to be part of the prompt (e.g. simple text + # in text out for Python) + list_tools_result.tools = [ + tool + for tool in list_tools_result.tools + if getattr(tool.annotations, "include_in_prompt", True) + ] + + return list_tools_result + + +class ToolServer(ABC): + + @abstractmethod + def has_tool(self, tool_name: str): + pass + + @abstractmethod + def get_tool_description(self, tool_name: str): + pass + + @abstractmethod + def get_tool_session(self, tool_name: str) -> AbstractAsyncContextManager[Any]: ... + + +class MCPToolServer(ToolServer): + + def __init__(self): + self.harmony_tool_descriptions = {} + + async def add_tool_server(self, server_url: str): + tool_urls = server_url.split(",") + self.harmony_tool_descriptions = {} + self.urls: dict[str, str] = {} + for url in tool_urls: + url = f"http://{url}/sse" + initialize_response, list_tools_response = await list_server_and_tools(url) + + list_tools_response = post_process_tools_description(list_tools_response) + + tool_from_mcp = ToolNamespaceConfig( + name=initialize_response.serverInfo.name, + description=initialize_response.instructions, + tools=[ + ToolDescription.new( + name=tool.name, + description=tool.description, + parameters=tool.inputSchema, + ) + for tool in list_tools_response.tools + ], + ) + self.harmony_tool_descriptions[tool_from_mcp.name] = tool_from_mcp + if tool_from_mcp.name not in self.urls: + self.urls[tool_from_mcp.name] = url + else: + logger.warning( + "Tool %s already exists. Ignoring duplicate tool server %s", + tool_from_mcp.name, + url, + ) + + def has_tool(self, tool_name: str): + return tool_name in self.harmony_tool_descriptions + + def get_tool_description(self, tool_name: str): + return self.harmony_tool_descriptions.get(tool_name) + + @asynccontextmanager + async def get_tool_session(self, tool_name: str): + url = self.urls.get(tool_name) + if url: + async with sse_client(url=url) as streams, ClientSession( + *streams + ) as session: + await session.initialize() + yield session + else: + logger.warning("Tool %s not found", tool_name) + + +class DemoToolServer(ToolServer): + + def __init__(self): + from sglang.srt.entrypoints.tool import ( + HarmonyBrowserTool, + HarmonyPythonTool, + Tool, + ) + + self.tools: dict[str, Tool] = {} + browser_tool = HarmonyBrowserTool() + if browser_tool.enabled: + self.tools["browser"] = browser_tool + python_tool = HarmonyPythonTool() + if python_tool.enabled: + self.tools["python"] = python_tool + + def has_tool(self, tool_name: str): + return tool_name in self.tools + + def get_tool_description(self, tool_name: str): + if tool_name not in self.tools: + return None + if tool_name == "browser": + return ToolNamespaceConfig.browser() + elif tool_name == "python": + return ToolNamespaceConfig.python() + else: + raise ValueError(f"Unknown tool {tool_name}") + + @asynccontextmanager + async def get_tool_session(self, tool_name: str): + yield self.tools[tool_name] diff --git a/python/sglang/srt/entrypoints/openai/usage_processor.py b/python/sglang/srt/entrypoints/openai/usage_processor.py new file mode 100644 index 000000000..cdd14d74e --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/usage_processor.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Mapping, Optional, final + +from sglang.srt.entrypoints.openai.protocol import UsageInfo + + +@final +class UsageProcessor: + """Stateless helpers that turn raw token counts into a UsageInfo.""" + + @staticmethod + def _details_if_cached(count: int) -> Optional[Dict[str, int]]: + """Return {"cached_tokens": N} only when N > 0 (keeps JSON slim).""" + return {"cached_tokens": count} if count > 0 else None + + @staticmethod + def calculate_response_usage( + responses: List[Dict[str, Any]], + n_choices: int = 1, + enable_cache_report: bool = False, + ) -> UsageInfo: + completion_tokens = sum(r["meta_info"]["completion_tokens"] for r in responses) + + prompt_tokens = sum( + responses[i]["meta_info"]["prompt_tokens"] + for i in range(0, len(responses), n_choices) + ) + + cached_details = None + if enable_cache_report: + cached_total = sum( + r["meta_info"].get("cached_tokens", 0) for r in responses + ) + cached_details = UsageProcessor._details_if_cached(cached_total) + + return UsageProcessor.calculate_token_usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cached_tokens=cached_details, + ) + + @staticmethod + def calculate_streaming_usage( + prompt_tokens: Mapping[int, int], + completion_tokens: Mapping[int, int], + cached_tokens: Mapping[int, int], + n_choices: int, + enable_cache_report: bool = False, + ) -> UsageInfo: + # index % n_choices == 0 marks the first choice of a prompt + total_prompt_tokens = sum( + tok for idx, tok in prompt_tokens.items() if idx % n_choices == 0 + ) + total_completion_tokens = sum(completion_tokens.values()) + + cached_details = ( + UsageProcessor._details_if_cached(sum(cached_tokens.values())) + if enable_cache_report + else None + ) + + return UsageProcessor.calculate_token_usage( + prompt_tokens=total_prompt_tokens, + completion_tokens=total_completion_tokens, + cached_tokens=cached_details, + ) + + @staticmethod + def calculate_token_usage( + prompt_tokens: int, + completion_tokens: int, + cached_tokens: Optional[Dict[str, int]] = None, + ) -> UsageInfo: + """Calculate token usage information""" + return UsageInfo( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=cached_tokens, + ) diff --git a/python/sglang/srt/entrypoints/openai/utils.py b/python/sglang/srt/entrypoints/openai/utils.py new file mode 100644 index 000000000..94ac5458d --- /dev/null +++ b/python/sglang/srt/entrypoints/openai/utils.py @@ -0,0 +1,72 @@ +import logging +from typing import Any, Dict, List, Optional, Union + +from sglang.srt.entrypoints.openai.protocol import ( + ChatCompletionRequest, + CompletionRequest, + LogProbs, +) + +logger = logging.getLogger(__name__) + + +def to_openai_style_logprobs( + input_token_logprobs=None, + output_token_logprobs=None, + input_top_logprobs=None, + output_top_logprobs=None, +): + ret_logprobs = LogProbs() + + def append_token_logprobs(token_logprobs): + for logprob, _, token_text in token_logprobs: + ret_logprobs.tokens.append(token_text) + ret_logprobs.token_logprobs.append(logprob) + + # Not supported yet + ret_logprobs.text_offset.append(-1) + + def append_top_logprobs(top_logprobs): + for tokens in top_logprobs: + if tokens is not None: + ret_logprobs.top_logprobs.append( + {token[2]: token[0] for token in tokens} + ) + else: + ret_logprobs.top_logprobs.append(None) + + if input_token_logprobs is not None: + append_token_logprobs(input_token_logprobs) + if output_token_logprobs is not None: + append_token_logprobs(output_token_logprobs) + if input_top_logprobs is not None: + append_top_logprobs(input_top_logprobs) + if output_top_logprobs is not None: + append_top_logprobs(output_top_logprobs) + + return ret_logprobs + + +def process_hidden_states_from_ret( + ret_item: Dict[str, Any], + request: Union[ + ChatCompletionRequest, + CompletionRequest, + ], +) -> Optional[List]: + """Process hidden states from a ret item in non-streaming response. + + Args: + ret_item: Response item containing meta_info + request: The original request object + + Returns: + Processed hidden states for the last token, or None + """ + if not request.return_hidden_states: + return None + + hidden_states = ret_item["meta_info"].get("hidden_states", None) + if hidden_states is not None: + hidden_states = hidden_states[-1] if len(hidden_states) > 1 else [] + return hidden_states diff --git a/python/sglang/srt/entrypoints/tool.py b/python/sglang/srt/entrypoints/tool.py new file mode 100644 index 000000000..45b87ac3a --- /dev/null +++ b/python/sglang/srt/entrypoints/tool.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +import logging +import os +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +from sglang.srt.utils import print_info_once, print_warning_once + +if TYPE_CHECKING: + # Avoid circular import. + from sglang.srt.entrypoints.context import ConversationContext + +logger = logging.getLogger(__name__) + + +class Tool(ABC): + + @abstractmethod + async def get_result(self, context: "ConversationContext") -> Any: + pass + + +class HarmonyBrowserTool(Tool): + + def __init__(self): + self.enabled = True + exa_api_key = os.getenv("EXA_API_KEY") + if not exa_api_key: + self.enabled = False + print_warning_once("EXA_API_KEY is not set, browsing is disabled") + return + + try: + from gpt_oss.tools.simple_browser import SimpleBrowserTool + from gpt_oss.tools.simple_browser.backend import ExaBackend + except ImportError: + self.enabled = False + print_warning_once("gpt_oss is not installed, browsing is disabled") + return + + browser_backend = ExaBackend(source="web", api_key=exa_api_key) + self.browser_tool = SimpleBrowserTool(backend=browser_backend) + print_info_once("Browser tool initialized") + + async def get_result(self, context: "ConversationContext") -> Any: + from sglang.srt.entrypoints.context import HarmonyContext + + assert isinstance(context, HarmonyContext) + last_msg = context.messages[-1] + tool_output_msgs = [] + async for msg in self.browser_tool.process(last_msg): + tool_output_msgs.append(msg) + return tool_output_msgs + + @property + def tool_config(self) -> Any: + return self.browser_tool.tool_config + + +class HarmonyPythonTool(Tool): + + def __init__(self): + self.enabled = True + + try: + from gpt_oss.tools.python_docker.docker_tool import PythonTool + except ImportError: + self.enabled = False + print_warning_once("gpt_oss is not installed, code interpreter is disabled") + return + + self.python_tool = PythonTool() + print_info_once("Code interpreter tool initialized") + + async def get_result(self, context: "ConversationContext") -> Any: + from sglang.srt.entrypoints.context import HarmonyContext + + assert isinstance(context, HarmonyContext) + last_msg = context.messages[-1] + tool_output_msgs = [] + async for msg in self.python_tool.process(last_msg): + tool_output_msgs.append(msg) + return tool_output_msgs + + @property + def tool_config(self) -> Any: + return self.python_tool.tool_config diff --git a/python/sglang/srt/eplb/__init__.py b/python/sglang/srt/eplb/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/eplb/eplb_algorithms/__init__.py b/python/sglang/srt/eplb/eplb_algorithms/__init__.py new file mode 100644 index 000000000..e2a267810 --- /dev/null +++ b/python/sglang/srt/eplb/eplb_algorithms/__init__.py @@ -0,0 +1,63 @@ +from enum import Enum, auto +from typing import Optional + +import torch + +from sglang.srt.eplb.eplb_algorithms import deepseek, deepseek_vec + + +class EplbAlgorithm(Enum): + deepseek = auto() + deepseek_hierarchical = auto() + deepseek_vec = auto() + deepseek_vec_hierarchical = auto() + # TODO may have more algorithm later + + +def rebalance_experts( + tokens_per_expert: torch.Tensor, + num_physical_experts: int, + num_local_physical_experts: int, + num_groups: Optional[int], + num_nodes: int, + algorithm: EplbAlgorithm, +): + if algorithm in [EplbAlgorithm.deepseek, EplbAlgorithm.deepseek_hierarchical]: + return deepseek.rebalance_experts( + weight=tokens_per_expert.sum(dim=0), + num_replicas=num_physical_experts, + num_groups=num_groups, + num_nodes=num_nodes, + num_gpus=num_physical_experts // num_local_physical_experts, + enable_hierarchical=algorithm == EplbAlgorithm.deepseek_hierarchical, + ) + + if algorithm in [ + EplbAlgorithm.deepseek_vec, + EplbAlgorithm.deepseek_vec_hierarchical, + ]: + return deepseek_vec.rebalance_experts( + tokens_per_expert=tokens_per_expert, + num_physical_experts=num_physical_experts, + num_local_physical_experts=num_local_physical_experts, + num_groups=num_groups, + num_nodes=num_nodes, + enable_hierarchical=algorithm == EplbAlgorithm.deepseek_vec_hierarchical, + ) + + raise NotImplementedError + + +def compute_algorithm( + raw_algorithm: str, + num_groups: Optional[int], + num_nodes: int, +) -> EplbAlgorithm: + if raw_algorithm != "auto": + return EplbAlgorithm[raw_algorithm] + + # TODO test on real scenarios and know which ones perform better + if (num_groups is not None) and (num_groups % num_nodes == 0): + return EplbAlgorithm.deepseek_hierarchical + else: + return EplbAlgorithm.deepseek diff --git a/python/sglang/srt/eplb/eplb_algorithms/deepseek.py b/python/sglang/srt/eplb/eplb_algorithms/deepseek.py new file mode 100644 index 000000000..180ccdee4 --- /dev/null +++ b/python/sglang/srt/eplb/eplb_algorithms/deepseek.py @@ -0,0 +1,223 @@ +# This file is copied from https://github.com/deepseek-ai/EPLB/blob/main/eplb.py since that one is not a pypi package +from typing import Tuple + +import torch + +from sglang.srt.utils import get_bool_env_var + + +def balanced_packing( + weight: torch.Tensor, num_packs: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Pack n weighted objects to m packs, such that each bin contains exactly n/m objects and the weights of all packs + are as balanced as possible. + + Parameters: + weight: [X, n], the weight of each item + num_packs: number of packs + + Returns: + pack_index: [X, n], the pack index of each item + rank_in_pack: [X, n], the rank of the item in the pack + """ + num_layers, num_groups = weight.shape + assert num_groups % num_packs == 0 + groups_per_pack = num_groups // num_packs + + if groups_per_pack == 1: + pack_index = torch.arange( + weight.size(-1), dtype=torch.int64, device=weight.device + ).expand(weight.shape) + rank_in_pack = torch.zeros_like(weight, dtype=torch.int64) + return pack_index, rank_in_pack + + indices = weight.float().sort(-1, descending=True).indices.cpu() + pack_index = torch.full_like(weight, fill_value=-1, dtype=torch.int64, device="cpu") + rank_in_pack = torch.full_like(pack_index, fill_value=-1) + for i in range(num_layers): + pack_weights = [0] * num_packs + pack_items = [0] * num_packs + for group in indices[i]: + pack = min( + (i for i in range(num_packs) if pack_items[i] < groups_per_pack), + key=pack_weights.__getitem__, + ) + assert pack_items[pack] < groups_per_pack + pack_index[i, group] = pack + rank_in_pack[i, group] = pack_items[pack] + pack_weights[pack] += weight[i, group] + pack_items[pack] += 1 + return pack_index, rank_in_pack + + +def replicate_experts( + weight: torch.Tensor, num_phy: int +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Replicate `num_log` experts to `num_phy` replicas, such that the maximum load of all replicas is minimized. + + Parameters: + weight: [X, num_log] + num_phy: total number of experts after replication + + Returns: + phy2log: [X, num_phy], logical expert id of each physical expert + rank: [X, num_phy], the replica rank + logcnt: [X, num_log], number of replicas for each logical expert + """ + n, num_log = weight.shape + num_redundant = num_phy - num_log + assert num_redundant >= 0 + device = weight.device + phy2log = torch.arange(num_phy, dtype=torch.int64, device=device).repeat(n, 1) + rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device) + logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device) + arangen = torch.arange(n, dtype=torch.int64, device=device) + for i in range(num_log, num_phy): + redundant_indices = (weight / logcnt).max(dim=-1).indices + phy2log[:, i] = redundant_indices + rank[:, i] = logcnt[arangen, redundant_indices] + logcnt[arangen, redundant_indices] += 1 + return phy2log, rank, logcnt + + +def rebalance_experts_hierarchical( + weight: torch.Tensor, + num_physical_experts: int, + num_groups: int, + num_nodes: int, + num_gpus: int, +): + """ + Parameters: + weight: [num_moe_layers, num_logical_experts] + num_physical_experts: number of physical experts after replication + num_groups: number of expert groups + num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster + num_gpus: number of GPUs, must be a multiple of `num_nodes` + + Returns: + physical_to_logical_map: [num_moe_layers, num_physical_experts] + logical_to_physical_map: [num_moe_layers, num_logical_experts, X] + logical_count: [num_moe_layers, num_logical_experts] + """ + num_layers, num_logical_experts = weight.shape + assert num_logical_experts % num_groups == 0 + group_size = num_logical_experts // num_groups + assert num_groups % num_nodes == 0 + groups_per_node = num_groups // num_nodes + assert num_gpus % num_nodes == 0 + assert num_physical_experts % num_gpus == 0 + phy_experts_per_gpu = num_physical_experts // num_gpus + + def inverse(perm: torch.Tensor) -> torch.Tensor: + inv = torch.empty_like(perm) + inv.scatter_( + 1, + perm, + torch.arange(perm.size(1), dtype=torch.int64, device=perm.device).expand( + perm.shape + ), + ) + return inv + + # Step 1: pack groups to nodes + tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1) + group_pack_index, group_rank_in_pack = balanced_packing(tokens_per_group, num_nodes) + log2mlog = ( + ( + (group_pack_index * groups_per_node + group_rank_in_pack) * group_size + ).unsqueeze(-1) + + torch.arange(group_size, dtype=torch.int64, device=group_pack_index.device) + ).flatten(-2) + mlog2log = inverse(log2mlog) + + # Step 2: construct redundant experts within nodes + # [num_layers * num_nodes, num_logical_experts // num_nodes] + tokens_per_mlog = weight.gather(-1, mlog2log).view( + -1, num_logical_experts // num_nodes + ) + phy2mlog, phyrank, mlogcnt = replicate_experts( + tokens_per_mlog, num_physical_experts // num_nodes + ) + + # Step 3: pack physical_experts to GPUs + # [num_layers * num_nodes, num_physical_experts // num_nodes] + tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog) + pack_index, rank_in_pack = balanced_packing(tokens_per_phy, num_gpus // num_nodes) + phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack + pphy2phy = inverse(phy2pphy) + + pphy2mlog = phy2mlog.gather( + -1, pphy2phy + ) # [num_layers * num_nodes, num_log_per_nodes] + pphy2mlog = ( + pphy2mlog.view(num_layers, num_nodes, -1) + + torch.arange( + 0, + num_logical_experts, + num_logical_experts // num_nodes, + device=group_pack_index.device, + ).view(1, -1, 1) + ).flatten(-2) + pphy2log = mlog2log.gather(-1, pphy2mlog) + pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1) + logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog) + return pphy2log, pphyrank, logcnt + + +def rebalance_experts( + weight: torch.Tensor, + num_replicas: int, + num_groups: int, + num_nodes: int, + num_gpus: int, + enable_hierarchical: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Entry point for expert-parallelism load balancer. + + Parameters: + weight: [layers, num_logical_experts], the load statistics for all logical experts + num_replicas: number of physical experts, must be a multiple of `num_gpus` + num_groups: number of expert groups + num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster + num_gpus: number of GPUs, must be a multiple of `num_nodes` + + Returns: + physical_to_logical_map: [layers, num_replicas], the expert index of each replica + logical_to_physical_map: [layers, num_logical_experts, X], the replica indices for each expert + expert_count: [layers, num_logical_experts], number of physical replicas for each logical expert + """ + + num_layers, num_logical_experts = weight.shape + weight = weight.float().cpu() + if enable_hierarchical: + # use hierarchical load-balance policy + phy2log, phyrank, logcnt = rebalance_experts_hierarchical( + weight, num_replicas, num_groups, num_nodes, num_gpus + ) + else: + # use global load-balance policy + phy2log, phyrank, logcnt = rebalance_experts_hierarchical( + weight, num_replicas, 1, 1, num_gpus + ) + maxlogcnt = logcnt.max().item() + log2phy: torch.Tensor = torch.full( + (num_layers, num_logical_experts, maxlogcnt), + -1, + dtype=torch.int64, + device=logcnt.device, + ) + log2phy.view(num_layers, -1).scatter_( + -1, + phy2log * maxlogcnt + phyrank, + torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand( + num_layers, -1 + ), + ) + return phy2log, log2phy, logcnt + + +__all__ = ["rebalance_experts"] diff --git a/python/sglang/srt/eplb/eplb_algorithms/deepseek_vec.py b/python/sglang/srt/eplb/eplb_algorithms/deepseek_vec.py new file mode 100644 index 000000000..cb165448a --- /dev/null +++ b/python/sglang/srt/eplb/eplb_algorithms/deepseek_vec.py @@ -0,0 +1,276 @@ +# This file is copied from https://github.com/deepseek-ai/EPLB/blob/main/eplb.py since that one is not a pypi package +from typing import Optional, Tuple + +import torch + + +def pack_groups(tokens_per_group: torch.Tensor, num_nodes: int) -> torch.Tensor: + num_layers, num_groups = tokens_per_group.shape + assert num_groups % num_nodes == 0 + groups_per_rank = num_groups // num_nodes + + indices = tokens_per_group.float().sort(-1, descending=True).indices.cpu() + ret = torch.full_like( + tokens_per_group, fill_value=-1, dtype=torch.int64, device="cpu" + ) + for layer in range(num_layers): + node_tokens = [0] * num_nodes + node_groups = [0] * num_nodes + for group in indices[layer]: + + def key_func(rank: int) -> int: + if node_groups[rank] >= groups_per_rank: + return 1, 0 + else: + return 0, node_tokens[rank] + + rank = min(range(num_nodes), key=key_func) + assert node_groups[rank] < groups_per_rank + ret[layer, group] = rank * groups_per_rank + node_groups[rank] + node_tokens[rank] += tokens_per_group[layer, group] + node_groups[rank] += 1 + return ret + + +def make_redundant_experts_chunkwise( + tokens_per_expert: torch.Tensor, + num_physical_experts: int, + num_local_physical_experts: int, + num_physical_experts_per_chunk: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + num_steps, num_moe_layers, num_logical_experts = tokens_per_expert.shape + num_redundancy_experts = num_physical_experts - num_logical_experts + + physical_to_logical_map = torch.empty( + num_moe_layers, + num_physical_experts, + dtype=torch.int, + device=tokens_per_expert.device, + ) + logical_to_physical_map = torch.full( + (num_moe_layers, num_logical_experts, num_redundancy_experts + 1), + -1, + dtype=torch.int, + device=tokens_per_expert.device, + ) + logical_count = torch.ones( + num_moe_layers, + num_logical_experts, + dtype=torch.int, + device=tokens_per_expert.device, + ) + + assert num_physical_experts % num_physical_experts_per_chunk == 0 + num_chunks = num_physical_experts // num_physical_experts_per_chunk + assert num_logical_experts % num_chunks == 0 + num_logical_experts_per_group = num_logical_experts // num_chunks + assert num_redundancy_experts % num_chunks == 0 + num_redundancy_experts_per_group = num_redundancy_experts // num_chunks + + arange_num_moe_layers_num_groups = torch.arange( + num_moe_layers * num_chunks, dtype=torch.int, device=tokens_per_expert.device + ) + arange_num_logical_experts = torch.arange( + num_logical_experts, dtype=torch.int, device=tokens_per_expert.device + ) + arange_num_logical_experts_per_group = torch.arange( + num_logical_experts_per_group, dtype=torch.int, device=tokens_per_expert.device + ) + arange_num_groups = torch.arange( + num_chunks, dtype=torch.int, device=tokens_per_expert.device + ) + physical_to_logical_map.view( + num_moe_layers, num_chunks, num_physical_experts_per_chunk + )[:, :, :num_logical_experts_per_group] = arange_num_logical_experts.view( + num_chunks, num_logical_experts_per_group + ) + logical_to_physical_map[:, :, 0] = ( + arange_num_logical_experts_per_group.expand( + num_chunks, num_logical_experts_per_group + ) + + arange_num_groups[:, None] * num_physical_experts_per_chunk + ).view(num_logical_experts) + + tokens_per_expert_all_diff = tokens_per_expert + arange_num_logical_experts * 1e-4 + for i in range(num_redundancy_experts_per_group): + score = ( + tokens_per_expert_all_diff / logical_count + ) # NOTE: Values in score must be different from each other + score1 = tokens_per_expert / (logical_count + 1) + score = score.view( + num_steps, num_moe_layers, num_chunks, num_logical_experts_per_group + ) + score1 = score1.view_as(score) + values, indices = score.max(-1, keepdim=True) + values = values.expand_as(score).contiguous() + score.scatter_(-1, indices, score1.gather(-1, indices)) + values.scatter_(-1, indices, score.max(-1, keepdim=True).values) + redundancy_indices = values.sum(0).argmin(-1) + physical_to_logical_map.view( + num_moe_layers, num_chunks, num_physical_experts_per_chunk + )[:, :, num_logical_experts_per_group + i] = ( + redundancy_indices + arange_num_groups * num_logical_experts_per_group + ) + redundancy_count = ( + logical_count.view( + num_moe_layers * num_chunks, num_logical_experts_per_group + ) + .gather(-1, redundancy_indices.view(num_moe_layers * num_chunks, 1)) + .squeeze(1) + ) + physical_redundancy_indices = ( + ( + arange_num_groups * num_physical_experts_per_chunk + + num_logical_experts_per_group + + i + ) + .expand(num_moe_layers, num_chunks) + .flatten() + ) + logical_to_physical_map.view( + num_moe_layers * num_chunks, + num_logical_experts_per_group, + num_redundancy_experts + 1, + )[ + arange_num_moe_layers_num_groups, + redundancy_indices.view(num_moe_layers * num_chunks), + redundancy_count, + ] = physical_redundancy_indices + logical_count.view(num_moe_layers * num_chunks, num_logical_experts_per_group)[ + arange_num_moe_layers_num_groups, + redundancy_indices.view(num_moe_layers * num_chunks), + ] += 1 + + if num_local_physical_experts > 1: + # Load-balancing between GPUs + physical_to_logical_map_int64 = physical_to_logical_map.to(torch.int64) + counts = logical_count.gather(-1, physical_to_logical_map_int64) + score = tokens_per_expert.sum(0).gather(-1, physical_to_logical_map_int64) + score = score / counts + score = score.view(num_moe_layers, num_chunks, num_physical_experts_per_chunk) + indices = score.argsort(-1, descending=True) + indices += torch.arange( + 0, + num_physical_experts, + num_physical_experts_per_chunk, + dtype=indices.dtype, + device=indices.device, + )[None, :, None] + + assert num_physical_experts_per_chunk % num_local_physical_experts == 0 + num_local_groups = num_physical_experts_per_chunk // num_local_physical_experts + indices = indices.view( + num_moe_layers, num_chunks, num_local_physical_experts, num_local_groups + ) + indices[:, :, 1::2, :] = indices[:, :, 1::2, :].flip(-1) + indices = indices.transpose(2, 3) + indices = indices.reshape(num_moe_layers, num_physical_experts) + physical_to_logical_map = physical_to_logical_map.gather(-1, indices) + mask = logical_to_physical_map == -1 + logical_to_physical_map[mask] = 0 + logical_to_physical_map = ( + indices.argsort(-1) + .gather( + -1, logical_to_physical_map.view(num_moe_layers, -1).to(torch.int64) + ) + .view_as(logical_to_physical_map) + .to(torch.int) + ) + logical_to_physical_map[mask] = -1 + + return physical_to_logical_map, logical_to_physical_map, logical_count + + +def decode_rebalance_experts( + tokens_per_expert: torch.Tensor, + num_physical_experts: int, + num_local_physical_experts: int, +): + return make_redundant_experts_chunkwise( + tokens_per_expert, + num_physical_experts, + num_local_physical_experts, + num_physical_experts, + ) + + +def prefill_rebalance_experts( + tokens_per_expert: torch.Tensor, + num_physical_experts: int, + num_local_physical_experts: int, + num_groups: int, + num_nodes: int, +): + tokens_per_expert = tokens_per_expert.float().cpu() + + num_steps, _, num_logical_experts = tokens_per_expert.shape + assert num_logical_experts % num_groups == 0 + group_size = num_logical_experts // num_groups + assert num_groups % num_nodes == 0, f"{num_groups=} {num_nodes=}" + + tokens_per_group = tokens_per_expert.sum(0).unflatten(-1, (num_groups, -1)).sum(-1) + group_perm = pack_groups( + tokens_per_group, num_nodes + ) # [num_moe_layers, num_groups] => [num_moe_layers, num_nodes] + + # log2mlog [layers, #logexp] -> [layers, #logexp] + log2mlog = ( + (group_perm * group_size).unsqueeze(-1) + + torch.arange(group_size, dtype=torch.int64, device=group_perm.device) + ).flatten(-2) + + # mlog2log [layers, #logexp] -> [layers, #logexp], inverse of log2mlog + mlog2log = torch.empty_like(log2mlog) + arange = torch.arange( + num_logical_experts, dtype=torch.int64, device=mlog2log.device + ) + mlog2log.scatter_(1, log2mlog, arange.expand(log2mlog.size(0), -1)) + + # tokens_per_mlog[i][j][k] = tokens_per_expert[i][j][mlog2log[j][k]] + tokens_per_mlog = tokens_per_expert.gather( + 2, mlog2log.unsqueeze(0).expand(num_steps, -1, -1) + ) + + phy2mlog, mlog2phy, mlog_count = make_redundant_experts_chunkwise( + tokens_per_mlog, + num_physical_experts, + num_local_physical_experts, + num_physical_experts // num_nodes, + ) + + # phy2log[i][j] = mlog2log[i][phy2mlog[i][j]] + phy2log = mlog2log.gather(1, phy2mlog.to(torch.int64)) + + # mlog2phy: [num_moe_layers, num_logical_experts, ...] + # log2phy[i][j][k] = mlog2phy[i][log2mlog[i][j]][k] + log2phy = mlog2phy.gather( + 1, log2mlog.unsqueeze(-1).expand(-1, -1, mlog2phy.size(-1)).to(torch.int64) + ) + + # log_count[i][j] = mlog_count[i][log2mlog[i][j]] + log_count = mlog_count.gather(1, log2mlog) + return phy2log, log2phy, log_count + + +def rebalance_experts( + tokens_per_expert: torch.Tensor, + num_physical_experts: int, + num_local_physical_experts: int, + num_groups: Optional[int], + num_nodes: int, + enable_hierarchical: bool, +): + if enable_hierarchical: + return prefill_rebalance_experts( + tokens_per_expert=tokens_per_expert, + num_physical_experts=num_physical_experts, + num_local_physical_experts=num_local_physical_experts, + num_groups=num_groups, + num_nodes=num_nodes, + ) + else: + return decode_rebalance_experts( + tokens_per_expert=tokens_per_expert, + num_physical_experts=num_physical_experts, + num_local_physical_experts=num_local_physical_experts, + ) diff --git a/python/sglang/srt/eplb/eplb_manager.py b/python/sglang/srt/eplb/eplb_manager.py new file mode 100644 index 000000000..e88a3d28e --- /dev/null +++ b/python/sglang/srt/eplb/eplb_manager.py @@ -0,0 +1,118 @@ +import logging +import time +from typing import TYPE_CHECKING, List + +import torch.cuda + +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location import ExpertLocationMetadata + +if TYPE_CHECKING: + from sglang.srt.model_executor.model_runner import ModelRunner + +logger = logging.getLogger(__name__) + + +class EPLBManager: + def __init__(self, model_runner: "ModelRunner"): + super().__init__() + self._model_runner = model_runner + self._server_args = model_runner.server_args + self._rebalance_layers_per_chunk = ( + self._server_args.eplb_rebalance_layers_per_chunk + ) + self._rebalance_num_iterations = self._server_args.eplb_rebalance_num_iterations + + # Otherwise, the circular buffer will contain stale data. If the case is needed, it can be implemented. + assert ( + self._server_args.eplb_rebalance_num_iterations + >= self._server_args.expert_distribution_recorder_buffer_size + ), "eplb_rebalance_num_iterations must be greater than expert_distribution_recorder_buffer_size" + + if not get_global_expert_distribution_recorder().recording: + get_global_expert_distribution_recorder().start_record() + + logger.info( + f"[EPLBManager] system started, will rebalance per {self._rebalance_num_iterations} iterations." + ) + + self._main_generator = self._entrypoint() + + def on_forward_pass_end(self): + next(self._main_generator) + + # can be more complex if needed + def _entrypoint(self): + while True: + for _ in range(self._rebalance_num_iterations): + yield + + yield from self.rebalance() + + def rebalance(self): + logger.info("[EPLBManager] rebalance start") + + enable_timing = self._rebalance_layers_per_chunk is None + + if enable_timing: + torch.get_device_module().synchronize() + time_start = time.time() + + dump_record_output = get_global_expert_distribution_recorder().dump_record( + output_mode="object" + ) + logical_count = dump_record_output["logical_count"] + average_utilization_rate_over_window = dump_record_output[ + "average_utilization_rate_over_window" + ] + + # Check whether rebalancing is needed + if not self._check_rebalance_needed(average_utilization_rate_over_window): + return + + expert_location_metadata = ExpertLocationMetadata.init_by_eplb( + self._server_args, self._model_runner.model_config, logical_count + ) + + update_layer_ids_chunks = self._compute_update_layer_ids_chunks() + for chunk_index, update_layer_ids in enumerate(update_layer_ids_chunks): + if len(update_layer_ids_chunks) > 1: + yield + self._model_runner.update_expert_location( + expert_location_metadata, + update_layer_ids=update_layer_ids, + ) + + msg = f"[EPLBManager] rebalance end" + if enable_timing: + torch.get_device_module().synchronize() + time_end = time.time() + msg += f" time={time_end - time_start:.3f}s" + logger.info(msg) + + def _check_rebalance_needed(self, average_utilization_rate_over_window): + if average_utilization_rate_over_window is None: + return True + + if ( + average_utilization_rate_over_window + > self._server_args.eplb_min_rebalancing_utilization_threshold + ): + logger.info( + f"[EPLBManager] Skipped ep rebalancing: current GPU utilization {average_utilization_rate_over_window:.2f} > minimum rebalance threshold {self._server_args.eplb_min_rebalancing_utilization_threshold:.2f}" + ) + return False + + return True + + def _compute_update_layer_ids_chunks(self) -> List[List[int]]: + all_layer_ids = sorted( + list(self._model_runner.model.routed_experts_weights_of_layer.keys()) + ) + chunk_size = self._rebalance_layers_per_chunk or 1000000 + return list(_chunk_list(all_layer_ids, chunk_size=chunk_size)) + + +def _chunk_list(items: List, chunk_size): + for start_index in range(0, len(items), chunk_size): + yield items[start_index : start_index + chunk_size] diff --git a/python/sglang/srt/eplb/eplb_simulator/__init__.py b/python/sglang/srt/eplb/eplb_simulator/__init__.py new file mode 100644 index 000000000..a1fcbdf00 --- /dev/null +++ b/python/sglang/srt/eplb/eplb_simulator/__init__.py @@ -0,0 +1 @@ +from . import reader diff --git a/python/sglang/srt/eplb/eplb_simulator/reader.py b/python/sglang/srt/eplb/eplb_simulator/reader.py new file mode 100644 index 000000000..97405c319 --- /dev/null +++ b/python/sglang/srt/eplb/eplb_simulator/reader.py @@ -0,0 +1,51 @@ +from collections import defaultdict +from pathlib import Path + +import torch +from tqdm import tqdm + +from sglang.srt.eplb.expert_distribution import ( + _convert_global_physical_count_to_logical_count, +) + +convert_global_physical_count_to_logical_count = ( + _convert_global_physical_count_to_logical_count +) + + +def read_mode_per_pass(dir_data: Path): + """Read data from ExpertDistributionRecorder when recorded with mode `per_pass`""" + + # gpc := global_physical_count + gpc_of_forward_pass_and_rank = defaultdict(lambda: defaultdict()) + for path in tqdm(list(dir_data.glob("*.pt"))): + data_pack = torch.load(path, weights_only=True) + last_physical_to_logical_map = data_pack["last_physical_to_logical_map"] + for record in data_pack["records"]: + forward_pass_id = record["forward_pass_id"] + rank = record["rank"] + assert ( + gpc_of_forward_pass_and_rank[forward_pass_id].get(rank) is None + ), f"Duplicated {forward_pass_id=} {rank=}" + gpc_of_forward_pass_and_rank[forward_pass_id][rank] = record[ + "global_physical_count" + ] + + forward_pass_ids = sorted(gpc_of_forward_pass_and_rank.keys()) + print(f"Make {forward_pass_ids=} into array") + + items = [] + for forward_pass_id, gpc_of_rank in sorted(gpc_of_forward_pass_and_rank.items()): + gpc_of_rank_tensor = torch.stack( + [gpc for rank, gpc in sorted(gpc_of_rank.items())] + ).sum(dim=0) + items.append(gpc_of_rank_tensor) + + gpc_of_forward_pass = torch.stack(items) + print(f"{gpc_of_forward_pass.shape=}") + + return dict( + global_physical_count_of_forward_pass=gpc_of_forward_pass, + last_physical_to_logical_map=last_physical_to_logical_map, + forward_pass_ids=forward_pass_ids, + ) diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py new file mode 100644 index 000000000..3faf981ef --- /dev/null +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -0,0 +1,966 @@ +# Copyright 2023-2024 SGLang 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 __future__ import annotations + +import logging +import math +import os +import time +from abc import ABC +from collections import deque +from contextlib import contextmanager +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Type + +import einops +import torch +import torch.distributed + +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import Withable, get_bool_env_var, is_npu + +_is_npu = is_npu() + +if TYPE_CHECKING: + from sglang.srt.eplb.expert_location import ExpertLocationMetadata + +logger = logging.getLogger(__name__) + +# --------------------------------------- Entrypoint ----------------------------------------- + +_OutputMode = Literal["file", "object"] + + +class ExpertDistributionRecorder(ABC): + """Global expert distribution recording""" + + @staticmethod + def init_new( + server_args: ServerArgs, + expert_location_metadata: ExpertLocationMetadata, + rank: int, + ): + if server_args.expert_distribution_recorder_mode is not None: + assert ( + expert_location_metadata is not None + ), "ExpertLocationMetadata is required for expert distribution recording. One possible" + "reason is that you are using a model that does not support expert distribution" + "recording. Try setting `get_model_config_for_expert_location` in your model." + return _ExpertDistributionRecorderReal( + server_args, expert_location_metadata, rank + ) + else: + return _ExpertDistributionRecorderNoop() + + @contextmanager + def with_current_layer(self, layer_idx): + yield + + @contextmanager + def with_debug_name(self, debug_name): + yield + + @contextmanager + def disable_this_region(self): + yield + + @contextmanager + def with_forward_pass(self, forward_pass_id: int, forward_batch: ForwardBatch): + yield + + def on_select_experts(self, topk_ids: torch.Tensor): + pass + + def on_deepep_dispatch_normal( + self, + local_physical_count_of_layer: List[int], + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + ): + pass + + def on_deepep_dispatch_low_latency( + self, local_physical_count_of_layer: torch.Tensor + ): + pass + + def start_record(self): + self._on_not_implemented() + + def stop_record(self): + self._on_not_implemented() + + def dump_record(self, output_mode: _OutputMode = "file"): + self._on_not_implemented() + + @property + def recording(self): + return False + + def _on_not_implemented(self): + raise Exception( + "Please set ServerArgs.expert_distribution_recorder_mode to use ExpertDistributionRecorder." + ) + + +class _ExpertDistributionRecorderNoop(ExpertDistributionRecorder): + pass + + +class _ExpertDistributionRecorderReal(ExpertDistributionRecorder): + def __init__( + self, + server_args: ServerArgs, + expert_location_metadata: ExpertLocationMetadata, + rank: int, + ): + self._server_args = server_args + self._expert_location_metadata = expert_location_metadata + + self._recording = False + self._disable_all = False + self._current_forward_pass_id = Withable() + self._current_layer_idx = Withable() + self._current_debug_name = Withable() + self._accumulator = _Accumulator.init_new( + server_args, expert_location_metadata, rank + ) + self._single_pass_gatherers = { + k: _SinglePassGatherer.init_new(server_args, expert_location_metadata, rank) + for k in self._accumulator.get_single_pass_gatherer_keys() + } + + if server_args.enable_expert_distribution_metrics: + logger.info( + "ExpertDistributionRecorder auto start record since enable_expert_distribution_metrics" + ) + self.start_record() + + def with_current_layer(self, layer_idx): + return self._current_layer_idx.with_value(layer_idx) + + def with_debug_name(self, debug_name): + return self._current_debug_name.with_value(debug_name) + + @contextmanager + def with_forward_pass(self, forward_pass_id: int, forward_batch: ForwardBatch): + with self._current_forward_pass_id.with_value(forward_pass_id): + self._on_forward_pass_start(forward_batch) + try: + yield + finally: + self._on_forward_pass_end(forward_pass_id) + + @contextmanager + def disable_this_region(self): + """Context manager to temporarily disable recording.""" + previous_disable_all = self._disable_all + self._disable_all = True + try: + yield + finally: + self._disable_all = previous_disable_all + + def _on_forward_pass_start(self, forward_batch: ForwardBatch): + if not self._recording: + return + for gatherer_key, gatherer in self._single_pass_gatherers.items(): + gatherer.reset() + gatherer.on_forward_pass_start(forward_batch) + + def _on_forward_pass_end(self, forward_pass_id: int): + if not self._recording: + return + for gatherer_key, gatherer in self._single_pass_gatherers.items(): + single_pass_data = gatherer.collect() + self._accumulator.append(forward_pass_id, gatherer_key, single_pass_data) + + def on_select_experts(self, topk_ids: torch.Tensor): + self._on_hook("on_select_experts", topk_ids=topk_ids) + + def on_deepep_dispatch_normal( + self, + local_physical_count_of_layer: List[int], + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + ): + self._on_hook( + "on_deepep_dispatch_normal", + local_physical_count_of_layer=local_physical_count_of_layer, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + num_tokens_per_expert=num_tokens_per_expert, + ) + + def on_deepep_dispatch_low_latency( + self, local_physical_count_of_layer: torch.Tensor + ): + self._on_hook( + "on_deepep_dispatch_low_latency", + local_physical_count_of_layer=local_physical_count_of_layer, + ) + + def _on_hook(self, hook_name: str, **kwargs): + if self._disable_all: + return + if not ( + self._recording or torch.get_device_module().is_current_stream_capturing() + ): + return + gatherer = self._single_pass_gatherers[ + self._accumulator.get_single_pass_gatherer_key( + self._current_debug_name.value + ) + ] + getattr(gatherer, hook_name)(layer_idx=self._current_layer_idx.value, **kwargs) + + def _reset(self): + """Reset the expert distribution recorder.""" + logger.info("Resetting ExpertDistributionRecorder...") + assert ( + self._current_layer_idx.value is None + ), f"{self._current_layer_idx.value=}" + for gatherer in self._single_pass_gatherers.values(): + gatherer.reset() + self._accumulator.reset() + + def start_record(self): + """Start recording the expert distribution.""" + if self._recording: + logger.warning( + "SGLang server is already recording expert ids. Did you forget to dump the expert ids recorded so far by sending requests to the `/stop_expert_distribution_record` and `/dump_expert_distribution_record` endpoints?" + ) + self._reset() + self._recording = True + + def stop_record(self): + """Stop recording the expert distribution.""" + if not self._recording: + logger.warning( + "SGLang server has not been recording expert ids. Did you forget to start recording by sending request to the `/start_expert_distribution_record` endpoint?" + ) + self._recording = False + + def dump_record(self, output_mode: _OutputMode = "file"): + """Dump the expert distribution record and reset the recorder after dumping.""" + output = self._accumulator.dump(output_mode=output_mode) + self._reset() + return output + + @property + def recording(self): + return self._recording + + +_global_expert_distribution_recorder: Optional[ExpertDistributionRecorder] = ( + _ExpertDistributionRecorderNoop() +) + + +def get_global_expert_distribution_recorder(): + return _global_expert_distribution_recorder + + +def set_global_expert_distribution_recorder(value): + global _global_expert_distribution_recorder + _global_expert_distribution_recorder = value + + +# --------------------------------------- SinglePassGatherer ----------------------------------------- + + +class _SinglePassGatherer(ABC): + @staticmethod + def init_new( + server_args: ServerArgs, + expert_location_metadata: ExpertLocationMetadata, + rank: int, + ) -> "_SinglePassGatherer": + if server_args.expert_distribution_recorder_mode == "per_token": + return _DetailSinglePassGatherer( + server_args, expert_location_metadata, rank + ) + + if server_args.expert_distribution_recorder_mode == "stat_approx": + if server_args.moe_a2a_backend != "none" and ( + server_args.deepep_mode == "normal" + ): + return _DeepepNormalSinglePassGatherer(expert_location_metadata, rank) + else: + raise NotImplementedError + + if server_args.moe_a2a_backend != "none": + if server_args.deepep_mode == "normal": + return _SelectExpertsSinglePassGatherer(expert_location_metadata, rank) + elif server_args.deepep_mode == "low_latency": + return _DeepepLowLatencySinglePassGatherer( + expert_location_metadata, rank + ) + else: + raise NotImplementedError + + return _SelectExpertsSinglePassGatherer(expert_location_metadata, rank) + + def __init__(self, expert_location_metadata: ExpertLocationMetadata, rank: int): + self._expert_location_metadata = expert_location_metadata + self._rank = rank + + def on_forward_pass_start(self, forward_batch: ForwardBatch): + pass + + def on_select_experts(self, layer_idx: int, topk_ids: torch.Tensor): + pass + + def on_deepep_dispatch_normal( + self, + layer_idx: int, + local_physical_count_of_layer: List[int], + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + ): + pass + + def on_deepep_dispatch_low_latency( + self, layer_idx: int, local_physical_count_of_layer: torch.Tensor + ): + pass + + def reset(self): + raise NotImplementedError + + def collect(self) -> Dict: + raise NotImplementedError + + +class _DetailSinglePassGatherer(_SinglePassGatherer): + # DeepSeek V3 has this value; should generalize later + _TOP_K_NUM = 8 + + def __init__( + self, + server_args: ServerArgs, + expert_location_metadata: ExpertLocationMetadata, + rank: int, + ): + super().__init__(expert_location_metadata, rank) + self._metadata: Optional[Dict[str, Any]] = None + self._topk_ids_of_layer = torch.zeros( + ( + expert_location_metadata.num_layers, + # TODO determine the max number + server_args.chunked_prefill_size * 8, + self._TOP_K_NUM, + ), + dtype=torch.int32, + device=server_args.device, + ) + self._misc_objects: List[Dict[str, Any]] = [] + assert ( + not server_args.enable_two_batch_overlap + ), "DetailSinglePassGatherer does not support TBO yet" + # TODO assert shared experts fusion is disabled, o/w data is wrong + + def on_forward_pass_start(self, forward_batch: ForwardBatch): + assert self._metadata is None + self._metadata = dict( + # TODO pr-chain + # rids=forward_batch.rids, + input_ids=forward_batch.input_ids.cpu().tolist(), + positions=forward_batch.positions.cpu().tolist(), + extend_seq_lens=forward_batch.extend_seq_lens_cpu, + forward_mode=forward_batch.forward_mode.value, + ) + + def on_select_experts(self, layer_idx: int, topk_ids: torch.Tensor): + self._topk_ids_of_layer[layer_idx, : topk_ids.shape[0], : topk_ids.shape[1]] = ( + topk_ids + ) + + def on_deepep_dispatch_normal( + self, + layer_idx: int, + local_physical_count_of_layer: List[int], + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + ): + self._misc_objects.append( + dict( + layer_id=layer_idx, + num_tokens_per_rank=num_tokens_per_rank.cpu().tolist(), + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank.cpu().tolist(), + num_tokens_per_expert=num_tokens_per_expert.cpu().tolist(), + ) + ) + + def reset(self): + self._topk_ids_of_layer[...] = -1 + self._misc_objects.clear() + self._metadata = None + + def collect(self) -> Dict: + num_tokens = len(self._metadata["input_ids"]) + return dict( + **self._metadata, + topk_ids_of_layer=self._topk_ids_of_layer[:, :num_tokens, :].clone().cpu(), + misc_objects=self._misc_objects, + ) + + +class _LayerBasedCpuSinglePassGatherer(_SinglePassGatherer): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._objects_of_layer = {} + + def _on_layer_data(self, layer_idx: int, objects: List[int]): + assert 0 <= layer_idx < self._expert_location_metadata.num_layers + if layer_idx in self._objects_of_layer: + self._objects_of_layer[layer_idx] = _list_sum( + self._objects_of_layer[layer_idx], objects + ) + else: + self._objects_of_layer[layer_idx] = objects + + def reset(self): + self._objects_of_layer.clear() + + def _collect_objects(self, pad_len: int) -> torch.Tensor: + data = [ + self._objects_of_layer.get(layer_index) or ([0] * pad_len) + for layer_index in range(self._expert_location_metadata.num_layers) + ] + return torch.tensor(data) + + +def _list_sum(a: List, b: List) -> List: + return [x + y for x, y in zip(a, b, strict=True)] + + +class _LayerBasedGpuSinglePassGatherer(_SinglePassGatherer): + def __init__(self, *args, enable_global_physical_experts: bool, **kwargs): + super().__init__(*args, **kwargs) + if not _is_npu: + device = "cuda" + else: + device = "npu" + self._enable_global_physical_experts = enable_global_physical_experts + self._data = torch.zeros( + ( + self._expert_location_metadata.num_layers, + ( + self._expert_location_metadata.num_physical_experts + if enable_global_physical_experts + else self._expert_location_metadata.num_local_physical_experts + ), + ), + dtype=torch.int, + device=device, + ) + + def reset(self): + self._data[...] = 0 + + def collect(self) -> Dict: + if self._enable_global_physical_experts: + global_physical_count = self._data + else: + # Can optimize if bottleneck + global_physical_count = _convert_local_to_global_physical_count( + self._data, + rank=self._rank, + num_local_physical_experts=self._expert_location_metadata.num_local_physical_experts, + num_physical_experts=self._expert_location_metadata.num_physical_experts, + ) + + return dict(global_physical_count=global_physical_count) + + +class _SelectExpertsSinglePassGatherer(_LayerBasedGpuSinglePassGatherer): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs, enable_global_physical_experts=True) + + # can optimize (e.g. fuse / compile) + def on_select_experts(self, layer_idx: int, topk_ids: torch.Tensor): + topk_ids = topk_ids.flatten() + mask = topk_ids != -1 + self._data[layer_idx, :].scatter_add_( + dim=0, index=topk_ids.masked_fill(~mask, 0).long(), src=mask.int() + ) + + +class _DeepepNormalSinglePassGatherer(_LayerBasedCpuSinglePassGatherer): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if torch.distributed.get_rank() == 0: + logger.info( + "DeepepNormalSinglePassGatherer gathers approximate statistics. " + "If used with small batch size, consider using expert_distribution_recorder_mode=stat." + ) + + def on_deepep_dispatch_normal( + self, + layer_idx: int, + local_physical_count_of_layer: List[int], + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + ): + assert isinstance(local_physical_count_of_layer, list) + self._on_layer_data(layer_idx, local_physical_count_of_layer) + + def collect(self) -> Dict: + local_physical_count = super()._collect_objects( + pad_len=self._expert_location_metadata.num_local_physical_experts + ) + global_physical_count = _convert_local_to_global_physical_count( + local_physical_count, + rank=self._rank, + num_local_physical_experts=self._expert_location_metadata.num_local_physical_experts, + num_physical_experts=self._expert_location_metadata.num_physical_experts, + ) + return dict(global_physical_count=global_physical_count) + + +class _DeepepLowLatencySinglePassGatherer(_LayerBasedGpuSinglePassGatherer): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs, enable_global_physical_experts=False) + + def on_deepep_dispatch_low_latency( + self, layer_idx: int, local_physical_count_of_layer: torch.Tensor + ): + # Most naive implementation, can optimize later + self._data[layer_idx, :] += local_physical_count_of_layer + + +def _convert_local_to_global_physical_count( + local_physical_count: torch.Tensor, + rank: int, + num_local_physical_experts: int, + num_physical_experts: int, +) -> torch.Tensor: + dtype = local_physical_count.dtype + device = local_physical_count.device + num_layers, _ = local_physical_count.shape + + ans = torch.zeros((num_layers, num_physical_experts), dtype=dtype, device=device) + ans[ + :, num_local_physical_experts * rank : num_local_physical_experts * (rank + 1) + ] = local_physical_count + return ans + + +# --------------------------------------- Accumulator ----------------------------------------- + +_SINGLE_PASS_GATHERER_KEY_PRIMARY = "primary" + + +class _Accumulator(ABC): + @staticmethod + def init_new( + server_args: ServerArgs, + expert_location_metadata: ExpertLocationMetadata, + rank: int, + ) -> "_Accumulator": + return _Accumulator.get_class(server_args)( + server_args, expert_location_metadata, rank + ) + + @staticmethod + def get_class(server_args: ServerArgs) -> Type["_Accumulator"]: + return { + "stat": _StatAccumulator, + "stat_approx": _StatAccumulator, + "per_pass": _DetailAccumulator, + "per_token": _DetailAccumulator, + }[server_args.expert_distribution_recorder_mode] + + def __init__( + self, + server_args: ServerArgs, + expert_location_metadata: ExpertLocationMetadata, + rank: int, + ): + self._server_args = server_args + self._expert_location_metadata = expert_location_metadata + self._rank = rank + + def get_single_pass_gatherer_keys(self): + return [_SINGLE_PASS_GATHERER_KEY_PRIMARY] + + def get_single_pass_gatherer_key(self, debug_name: Optional[str]): + return _SINGLE_PASS_GATHERER_KEY_PRIMARY + + def append( + self, + forward_pass_id: int, + gatherer_key: str, + single_pass_data: Dict, + ): + pass + + def reset(self): + pass + + def dump(self, output_mode: _OutputMode): + pass + + +class _UtilizationRateAccumulatorMixin(_Accumulator): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self._enable = self._server_args.enable_expert_distribution_metrics + + if self._enable: + self.window_sizes = [10, 100, 1000] + self._history = _DequeCollection(maxlens=self.window_sizes) + self._rank = torch.distributed.get_rank() + + def append( + self, + forward_pass_id: int, + gatherer_key: str, + single_pass_data: Dict, + ): + super().append(forward_pass_id, gatherer_key, single_pass_data) + if self._enable: + self._append_utilization_rate( + forward_pass_id, single_pass_data["global_physical_count"] + ) + + def reset(self): + super().reset() + if self._enable: + self._history.clear() + + def _append_utilization_rate( + self, forward_pass_id: int, single_pass_global_physical_count: torch.Tensor + ): + gpu_physical_count = compute_gpu_physical_count( + single_pass_global_physical_count, + num_gpu=self._expert_location_metadata.ep_size, + ) + gpu_physical_count = gpu_physical_count.to(self._server_args.device) + torch.distributed.reduce( + gpu_physical_count, dst=0, op=torch.distributed.ReduceOp.SUM + ) + + if self._rank == 0: + utilization_rate_tensor = compute_utilization_rate(gpu_physical_count) + utilization_rate = torch.mean(utilization_rate_tensor).item() + self._history.append(utilization_rate) + + gpu_physical_count_sum = gpu_physical_count.sum().item() + + logger.info( + f"[Expert Balancedness] " + f"forward_pass_id={forward_pass_id} " + f"current_pass_balancedness={utilization_rate:.03f} " + f"{''.join(f'last_{size}_average_balancedness={value:.03f} ' for size, value in self._history.mean().items())} " + f"gpu_physical_count_sum={gpu_physical_count_sum}" + # f"current_pass_per_layer={[round(x, 2) for x in utilization_rate_tensor.cpu().tolist()]}" + ) + + +class _DequeCollection: + def __init__(self, maxlens: List[int]): + self._dequeues = [deque(maxlen=maxlen) for maxlen in maxlens] + + def append(self, value): + for d in self._dequeues: + d.append(value) + + def clear(self): + for d in self._dequeues: + d.clear() + + def mean(self) -> Dict[int, float]: + return {d.maxlen: sum(d) / len(d) for d in self._dequeues} + + +class _DetailAccumulator(_UtilizationRateAccumulatorMixin): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._records = [] + + def get_single_pass_gatherer_keys(self): + if False: # TODO `server_args.enable_two_batch_overlap` + return [_SINGLE_PASS_GATHERER_KEY_PRIMARY, "child_a", "child_b"] + return super().get_single_pass_gatherer_keys() + + def get_single_pass_gatherer_key(self, debug_name: Optional[str]): + if False: # TODO `server_args.enable_two_batch_overlap` + return debug_name or _SINGLE_PASS_GATHERER_KEY_PRIMARY + return super().get_single_pass_gatherer_key(debug_name) + + def append( + self, + forward_pass_id: int, + gatherer_key: str, + single_pass_data: Dict, + ): + super().append(forward_pass_id, gatherer_key, single_pass_data) + + def _process_object(obj): + if isinstance(obj, torch.Tensor): + return obj.cpu().clone() + return obj + + single_pass_data_processed = { + k: _process_object(v) for k, v in single_pass_data.items() + } + + self._records.append( + dict( + forward_pass_id=forward_pass_id, + rank=self._rank, + gatherer_key=gatherer_key, + **single_pass_data_processed, + ) + ) + + def reset(self): + super().reset() + self._records.clear() + + def dump(self, output_mode: _OutputMode): + assert output_mode == "file" + output = dict( + records=self._records, + # NOTE: This may change during recording, so here we say it is the "last" one + last_physical_to_logical_map=self._expert_location_metadata.physical_to_logical_map, + ) + _dump_to_file( + f"expert_distribution_recorder_{time.time()}_{self._rank}.pt", output + ) + + +class _StatAccumulator(_UtilizationRateAccumulatorMixin): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._global_physical_count_of_buffered_step = _Buffer.init_new( + item_shape=( + self._expert_location_metadata.num_layers, + # Cannot use local_physical_count to support select_experts + self._expert_location_metadata.num_physical_experts, + ), + buffer_size=self._server_args.expert_distribution_recorder_buffer_size, + dtype=torch.int32, + device=self._server_args.device, + ) + self._first_dump = True + + def append( + self, + forward_pass_id: int, + gatherer_key: str, + single_pass_data: Dict, + ): + super().append(forward_pass_id, gatherer_key, single_pass_data) + # Can optimize if overhead here is large + self._global_physical_count_of_buffered_step.append( + single_pass_data["global_physical_count"] + ) + + def reset(self): + super().reset() + self._global_physical_count_of_buffered_step.reset() + + def dump(self, output_mode: _OutputMode): + logical_count_of_buffered_step = _convert_global_physical_count_to_logical_count( + self._global_physical_count_of_buffered_step.get_all(), + num_layers=self._expert_location_metadata.num_layers, + num_logical_experts=self._expert_location_metadata.num_logical_experts, + physical_to_logical_map=self._expert_location_metadata.physical_to_logical_map, + ) + + if self._first_dump: + self._first_dump = False + torch.get_device_module().empty_cache() + + torch.distributed.all_reduce( + logical_count_of_buffered_step, op=torch.distributed.ReduceOp.SUM + ) + + output = dict( + rank=self._rank, + logical_count=logical_count_of_buffered_step, + average_utilization_rate_over_window=self._get_global_average_utilization_rate(), + ) + + if output_mode == "file": + if self._rank == 0: + _dump_to_file(f"expert_distribution_recorder_{time.time()}.pt", output) + elif output_mode == "object": + return output + else: + raise NotImplementedError + + def _get_global_average_utilization_rate(self): + if not self._enable or math.isclose( + self._server_args.eplb_min_rebalancing_utilization_threshold, 1.0 + ): + return None + + if self._rank == 0: + utilization_mean_rates = self._history.mean() + window_index = self.window_sizes[-1] + average_utilization_rate_over_window = ( + utilization_mean_rates[window_index] + if window_index in utilization_mean_rates + else 0 + ) + + avg_rate_tensor = torch.tensor( + [average_utilization_rate_over_window], + dtype=torch.float32, + device="cuda", + ) + else: + avg_rate_tensor = torch.empty(1, dtype=torch.float32, device="cuda") + torch.distributed.broadcast(avg_rate_tensor, src=0) + return avg_rate_tensor.item() + + +def _dump_to_file(name, data): + save_dir = Path(os.environ.get("SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR", "/tmp")) + path_output = save_dir / name + logger.info(f"Write expert distribution to {path_output}") + if not save_dir.exists(): + save_dir.mkdir(parents=True, exist_ok=True) + torch.save(data, str(path_output)) + + +class _Buffer: + @staticmethod + def init_new(item_shape: Tuple, buffer_size: int, dtype, device): + if buffer_size < 0: + return _InfiniteBuffer(item_shape, dtype=dtype, device=device) + else: + return _CircularBuffer(item_shape, buffer_size, dtype=dtype, device=device) + + def append(self, value: torch.Tensor): + raise NotImplementedError + + def get_all(self) -> torch.Tensor: + raise NotImplementedError + + def reset(self): + raise NotImplementedError + + +class _CircularBuffer(_Buffer): + def __init__(self, item_shape: Tuple, buffer_size: int, dtype, device): + self._buffer = torch.zeros( + (buffer_size, *item_shape), dtype=dtype, device=device + ) + self._curr_index = 0 + + def append(self, value: torch.Tensor): + self._buffer[self._curr_index] = value + self._curr_index = (self._curr_index + 1) % len(self._buffer) + + def get_all(self) -> torch.Tensor: + return self._buffer + + def reset(self): + self._buffer[...] = 0 + + +class _InfiniteBuffer(_Buffer): + def __init__(self, item_shape: Tuple, dtype, device): + self._item_shape = item_shape + self._buffer = torch.zeros((128, *item_shape), dtype=dtype, device=device) + self._size = 0 + + def append(self, value: torch.Tensor): + curr_buffer_size = len(self._buffer) + dtype = self._buffer.dtype + device = self._buffer.device + + if self._size == curr_buffer_size: + new_buffer = torch.zeros( + (2 * curr_buffer_size, *self._item_shape), dtype=dtype, device=device + ) + new_buffer[:curr_buffer_size] = self._buffer + self._buffer = new_buffer + + self._buffer[self._size] = value + self._size += 1 + + def get_all(self) -> torch.Tensor: + return self._buffer[: self._size] + + def reset(self): + self._buffer[...] = 0 + self._size = 0 + + +def _convert_global_physical_count_to_logical_count( + # (whatever, num_layers, num_physical_experts) + global_physical_count: torch.Tensor, + num_layers: int, + num_logical_experts: int, + physical_to_logical_map: torch.Tensor, +): + dim_extra, _, _ = global_physical_count.shape + dtype = global_physical_count.dtype + device = global_physical_count.device + logical_count = torch.zeros( + (dim_extra, num_layers, num_logical_experts), dtype=dtype, device=device + ) + logical_count.scatter_add_( + dim=2, + index=physical_to_logical_map.unsqueeze(0) + .expand(dim_extra, -1, -1) + .to(torch.int64), + src=global_physical_count, + ) + return logical_count + + +def compute_gpu_physical_count( + physical_count_of_whatever: torch.Tensor, # (..., num_layer, num_physical_expert) + num_gpu: int, +): + """output: gpu_physical_count_of_batch (..., num_layer, num_gpu)""" + return einops.reduce( + physical_count_of_whatever, + "... num_layer (num_gpu num_expert_per_gpu) -> ... num_layer num_gpu", + "sum", + num_gpu=num_gpu, + ) + + +def compute_utilization_rate( + gpu_physical_count_of_batch: torch.Tensor, # (..., num_layer, num_gpu) +): + """output: utilization_rate (..., num_layer)""" + gpu_physical_count_of_batch = gpu_physical_count_of_batch.float() + max_gpu_physical_count = einops.reduce( + gpu_physical_count_of_batch, + "... num_layer num_gpu -> ... num_layer", + "max", + ) + avg_gpu_physical_count = einops.reduce( + gpu_physical_count_of_batch, + "... num_layer num_gpu -> ... num_layer", + "mean", + ) + return (avg_gpu_physical_count + 1e-5) / (max_gpu_physical_count + 1e-5) diff --git a/python/sglang/srt/eplb/expert_location.py b/python/sglang/srt/eplb/expert_location.py new file mode 100644 index 000000000..ee5f2c7ca --- /dev/null +++ b/python/sglang/srt/eplb/expert_location.py @@ -0,0 +1,468 @@ +# Copyright 2023-2024 SGLang 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 __future__ import annotations + +import json +import logging +import random +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, List, Optional + +import torch +import torch.distributed +import torch.nn.functional as F + +from sglang.srt.eplb import eplb_algorithms +from sglang.srt.model_loader import get_model_architecture + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +@dataclass +class ExpertLocationMetadata: + physical_to_logical_map: torch.Tensor # (layers, num_physical_experts) + physical_to_logical_map_cpu: torch.Tensor + logical_to_all_physical_map: torch.Tensor # (layers, num_logical_experts, X) + logical_to_all_physical_map_cpu: torch.Tensor # CPU copy for performance + logical_to_all_physical_map_num_valid: torch.Tensor # (layers, num_logical_experts) + # (layers, num_logical_experts) + logical_to_rank_dispatch_physical_map: Optional[torch.Tensor] + + # -------------------------------- properties ------------------------------------ + + @property + def num_layers(self) -> int: + return self.physical_to_logical_map.shape[0] + + @property + def num_physical_experts(self) -> int: + return self.physical_to_logical_map.shape[1] + + @property + def num_local_physical_experts(self) -> int: + ans, remainder = divmod(self.num_physical_experts, self.ep_size) + assert remainder == 0 + return ans + + @property + def num_logical_experts(self) -> int: + return self.logical_to_all_physical_map.shape[1] + + @property + def ep_size(self): + # TODO change when EP size != world size + return torch.distributed.get_world_size() + + def __post_init__(self): + num_layers_0, num_physical_experts_0 = self.physical_to_logical_map.shape + num_layers_1, num_logical_experts_0, num_physical_experts_1 = ( + self.logical_to_all_physical_map.shape + ) + num_layers_2, num_logical_experts_1 = ( + self.logical_to_all_physical_map_num_valid.shape + ) + assert num_layers_0 == num_layers_1 == num_layers_2 + assert num_logical_experts_0 == num_logical_experts_1 + assert num_physical_experts_0 == num_physical_experts_1 + + # -------------------------------- construction ------------------------------------ + + @staticmethod + def init_trivial(server_args: ServerArgs, model_config: ModelConfig): + """Trivial location - logical expert i corresponds to physical expert i""" + common = ExpertLocationMetadata._init_common(server_args, model_config) + + if common is None: + return None + + num_physical_experts = common["num_physical_experts"] + model_config_for_expert_location = common["model_config_for_expert_location"] + num_layers = model_config_for_expert_location.num_layers + num_logical_experts = model_config_for_expert_location.num_logical_experts + + physical_to_logical_map = ( + torch.arange(0, num_physical_experts).repeat(num_layers, 1) + % num_logical_experts + ) + + return ExpertLocationMetadata.init_by_mapping( + server_args, + model_config, + physical_to_logical_map=physical_to_logical_map, + ) + + @staticmethod + def init_by_mapping( + server_args: ServerArgs, + model_config: ModelConfig, + physical_to_logical_map, + ): + if not isinstance(physical_to_logical_map, torch.Tensor): + physical_to_logical_map = torch.tensor(physical_to_logical_map) + physical_to_logical_map = physical_to_logical_map.to(server_args.device) + + common = ExpertLocationMetadata._init_common(server_args, model_config) + + if common is None: + return None + + model_config_for_expert_location = common["model_config_for_expert_location"] + logical_to_all_physical_map = _compute_logical_to_all_physical_map( + physical_to_logical_map, + num_logical_experts=model_config_for_expert_location.num_logical_experts, + ) + + return ExpertLocationMetadata._init_raw( + server_args=server_args, + ep_size=common["ep_size"], + physical_to_logical_map=physical_to_logical_map, + logical_to_all_physical_map=logical_to_all_physical_map, + ) + + @staticmethod + def init_by_eplb( + server_args: ServerArgs, model_config: ModelConfig, logical_count: torch.Tensor + ): + if not isinstance(logical_count, torch.Tensor): + logical_count = torch.tensor(logical_count) + if len(logical_count.shape) == 2: + logical_count = logical_count.unsqueeze(0) + logical_count = logical_count.to(server_args.device) + + common = ExpertLocationMetadata._init_common(server_args, model_config) + + if common is None: + return None + + model_config_for_expert_location = common["model_config_for_expert_location"] + num_physical_experts = common["num_physical_experts"] + num_groups = model_config_for_expert_location.num_groups + num_nodes = server_args.nnodes + + physical_to_logical_map, logical_to_all_physical_map, expert_count = ( + eplb_algorithms.rebalance_experts( + tokens_per_expert=logical_count, + num_physical_experts=num_physical_experts, + num_local_physical_experts=num_physical_experts // common["ep_size"], + num_groups=num_groups, + num_nodes=num_nodes, + algorithm=eplb_algorithms.compute_algorithm( + raw_algorithm=server_args.eplb_algorithm, + num_groups=num_groups, + num_nodes=num_nodes, + ), + ) + ) + + return ExpertLocationMetadata._init_raw( + server_args=server_args, + ep_size=common["ep_size"], + physical_to_logical_map=physical_to_logical_map.to(server_args.device), + logical_to_all_physical_map=logical_to_all_physical_map.to( + server_args.device + ), + ) + + @staticmethod + def _init_common(server_args: ServerArgs, model_config: ModelConfig): + model_config_for_expert_location = ( + ModelConfigForExpertLocation.from_model_config(model_config) + ) + + if model_config_for_expert_location is None: + return None + + num_physical_experts = ( + model_config_for_expert_location.num_logical_experts + + server_args.ep_num_redundant_experts + ) + ep_size = server_args.ep_size + assert num_physical_experts % ep_size == 0 + num_local_physical_experts = num_physical_experts // ep_size + + return dict( + model_config_for_expert_location=model_config_for_expert_location, + num_physical_experts=num_physical_experts, + num_local_physical_experts=num_local_physical_experts, + ep_size=ep_size, + ) + + @staticmethod + def _init_raw( + server_args: ServerArgs, + ep_size: int, + physical_to_logical_map: torch.Tensor, + logical_to_all_physical_map: torch.Tensor, + ): + _, num_physical_experts = physical_to_logical_map.shape + + logical_to_all_physical_map_padded = F.pad( + logical_to_all_physical_map, + (0, num_physical_experts - logical_to_all_physical_map.shape[-1]), + value=-1, + ) + + logical_to_all_physical_map_num_valid = torch.count_nonzero( + logical_to_all_physical_map != -1, dim=-1 + ) + + return ExpertLocationMetadata( + physical_to_logical_map=physical_to_logical_map, + physical_to_logical_map_cpu=physical_to_logical_map.cpu(), + logical_to_all_physical_map=logical_to_all_physical_map_padded, + logical_to_all_physical_map_cpu=logical_to_all_physical_map_padded.cpu(), + logical_to_all_physical_map_num_valid=logical_to_all_physical_map_num_valid, + logical_to_rank_dispatch_physical_map=( + compute_logical_to_rank_dispatch_physical_map( + logical_to_all_physical_map=logical_to_all_physical_map, + num_gpus=ep_size, + num_physical_experts=num_physical_experts, + # TODO improve when we have real EP rank + ep_rank=torch.distributed.get_rank() % ep_size, + ) + if server_args.ep_dispatch_algorithm == "static" + else None + ), + ) + + # -------------------------------- mutation ------------------------------------ + + def update( + self, + other: "ExpertLocationMetadata", + update_layer_ids: List[int], + ): + for field in [ + "ep_size", + ]: + assert getattr(self, field) == getattr(other, field) + + for field in [ + "physical_to_logical_map", + "physical_to_logical_map_cpu", + "logical_to_all_physical_map", + "logical_to_all_physical_map_cpu", + "logical_to_all_physical_map_num_valid", + "logical_to_rank_dispatch_physical_map", + ]: + other_field = getattr(other, field) + self_field = getattr(self, field) + assert (other_field is not None) == (self_field is not None) + if self_field is not None: + mask_update = torch.tensor( + [i in update_layer_ids for i in range(self.num_layers)] + ) + mask_update = mask_update.view(*([-1] + [1] * (self_field.dim() - 1))) + mask_update = mask_update.to(self_field.device, non_blocking=True) + self_field[...] = torch.where(mask_update, other_field, self_field) + + # -------------------------------- usage ------------------------------------ + + def logical_to_all_physical( + self, layer_id: int, logical_expert_id: int + ) -> List[int]: + # Use CPU copy to avoid GPU→CPU sync on every call, which is expensive in update weights scenario + return [ + physical_expert_id + for physical_expert_id in self.logical_to_all_physical_map_cpu[ + layer_id, logical_expert_id + ].tolist() + if physical_expert_id != -1 + ] + + +_global_expert_location_metadata: Optional[ExpertLocationMetadata] = None + + +def get_global_expert_location_metadata(): + return _global_expert_location_metadata + + +def set_global_expert_location_metadata(value): + global _global_expert_location_metadata + assert _global_expert_location_metadata is None + _global_expert_location_metadata = value + + +def _compute_logical_to_all_physical_map( + physical_to_logical_map: torch.Tensor, num_logical_experts: int +): + # This is rarely called, so we use for loops for maximum clarity + + num_layers, num_physical_experts = physical_to_logical_map.shape + + logical_to_all_physical_map = [ + [[] for _ in range(num_logical_experts)] for _ in range(num_layers) + ] + for layer_id in range(num_layers): + for physical_expert_id in range(num_physical_experts): + logical_expert_id = physical_to_logical_map[ + layer_id, physical_expert_id + ].item() + logical_to_all_physical_map[layer_id][logical_expert_id].append( + physical_expert_id + ) + + logical_to_all_physical_map = _pad_nested_array( + logical_to_all_physical_map, pad_value=-1 + ) + + return torch.tensor( + logical_to_all_physical_map, device=physical_to_logical_map.device + ) + + +def _pad_nested_array(arr, pad_value): + max_len = max(len(inner) for outer in arr for inner in outer) + padded = [ + [inner + [pad_value] * (max_len - len(inner)) for inner in outer] + for outer in arr + ] + return padded + + +# TODO optimize performance (rewrite and/or run in separate process with overlap) +def compute_logical_to_rank_dispatch_physical_map( + logical_to_all_physical_map: torch.Tensor, + num_gpus: int, + num_physical_experts: int, + ep_rank: int, + seed: int = 42, +): + r = random.Random(seed) + + num_local_physical_experts = num_physical_experts // num_gpus + num_layers, num_logical_experts, _ = logical_to_all_physical_map.shape + dtype = logical_to_all_physical_map.dtype + + logical_to_rank_dispatch_physical_map = torch.full( + size=(num_gpus, num_layers, num_logical_experts), + fill_value=-1, + dtype=dtype, + ) + + for layer_id in range(num_layers): + for logical_expert_id in range(num_logical_experts): + candidate_physical_expert_ids = _logical_to_all_physical_raw( + logical_to_all_physical_map, layer_id, logical_expert_id + ) + output_partial = logical_to_rank_dispatch_physical_map[ + :, layer_id, logical_expert_id + ] + + for gpu_id in range(num_gpus): + same_gpu_physical_expert_ids = [ + physical_expert_id + for physical_expert_id in candidate_physical_expert_ids + if _compute_gpu_id_of_physical_expert( + physical_expert_id, num_local_physical_experts + ) + == gpu_id + ] + if len(same_gpu_physical_expert_ids) > 0: + output_partial[gpu_id] = same_gpu_physical_expert_ids[0] + + num_remain = torch.sum(output_partial == -1).item() + output_partial[output_partial == -1] = torch.tensor( + _fair_choices(candidate_physical_expert_ids, k=num_remain, r=r), + dtype=dtype, + ) + + assert torch.all(logical_to_rank_dispatch_physical_map != -1) + + device = logical_to_all_physical_map.device + return logical_to_rank_dispatch_physical_map[ep_rank, :, :].to(device) + + +def _logical_to_all_physical_raw( + logical_to_all_physical_map, layer_id: int, logical_expert_id: int +) -> List[int]: + return [ + physical_expert_id + for physical_expert_id in logical_to_all_physical_map[ + layer_id, logical_expert_id + ].tolist() + if physical_expert_id != -1 + ] + + +def _compute_gpu_id_of_physical_expert( + physical_expert_id: int, num_local_physical_experts: int +) -> int: + return physical_expert_id // num_local_physical_experts + + +def _fair_choices(arr: List, k: int, r: random.Random) -> List: + quotient, remainder = divmod(k, len(arr)) + ans = arr * quotient + r.sample(arr, k=remainder) + r.shuffle(ans) + return ans + + +@dataclass +class ModelConfigForExpertLocation: + num_layers: int + num_logical_experts: int + num_groups: Optional[int] = None + + @staticmethod + def from_model_config(model_config: ModelConfig): + model_class, _ = get_model_architecture(model_config) + if hasattr(model_class, "get_model_config_for_expert_location"): + return model_class.get_model_config_for_expert_location( + model_config.hf_config + ) + else: + return None + + +def compute_initial_expert_location_metadata( + server_args: ServerArgs, model_config: ModelConfig +) -> Optional[ExpertLocationMetadata]: + data = server_args.init_expert_location + if data == "trivial": + return ExpertLocationMetadata.init_trivial(server_args, model_config) + + # TODO unify with the utils function + if data.endswith(".pt"): + data_dict = torch.load(data, weights_only=True) + elif data.endswith(".json"): + data_dict = json.loads(Path(data).read_text()) + else: + data_dict = json.loads(data) + + if "physical_to_logical_map" in data_dict: + logger.info( + "init_expert_location from init_by_mapping using ServerArgs.init_expert_location" + ) + return ExpertLocationMetadata.init_by_mapping( + server_args, model_config, **data_dict + ) + elif "logical_count" in data_dict: + logger.info( + "init_expert_location from init_by_eplb using ServerArgs.init_expert_location" + ) + return ExpertLocationMetadata.init_by_eplb( + server_args, model_config, logical_count=data_dict["logical_count"] + ) + else: + raise NotImplementedError( + f"Unknown init_expert_location format ({list(data_dict.keys())=})" + ) diff --git a/python/sglang/srt/eplb/expert_location_dispatch.py b/python/sglang/srt/eplb/expert_location_dispatch.py new file mode 100644 index 000000000..624dc3fd1 --- /dev/null +++ b/python/sglang/srt/eplb/expert_location_dispatch.py @@ -0,0 +1,109 @@ +# Copyright 2023-2025 SGLang 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 dataclasses import dataclass +from typing import Literal, Optional + +import torch + +from sglang.srt.eplb.expert_location import get_global_expert_location_metadata +from sglang.srt.managers.schedule_batch import global_server_args_dict + + +@dataclass +class ExpertLocationDispatchInfo: + ep_dispatch_algorithm: Literal["static", "random"] + # (num_logical_experts,) + partial_logical_to_rank_dispatch_physical_map: Optional[torch.Tensor] + # (num_logical_experts, X) + partial_logical_to_all_physical_map: torch.Tensor + # (num_logical_experts,) + partial_logical_to_all_physical_map_num_valid: torch.Tensor + num_physical_experts: int + + @classmethod + def init_new(cls, layer_id: int): + ep_dispatch_algorithm = global_server_args_dict["ep_dispatch_algorithm"] + expert_location_metadata = get_global_expert_location_metadata() + assert expert_location_metadata is not None + + if ep_dispatch_algorithm is None: + return None + + return cls( + ep_dispatch_algorithm=ep_dispatch_algorithm, + partial_logical_to_rank_dispatch_physical_map=( + expert_location_metadata.logical_to_rank_dispatch_physical_map[ + layer_id, : + ] + if expert_location_metadata.logical_to_rank_dispatch_physical_map + is not None + else None + ), + partial_logical_to_all_physical_map=expert_location_metadata.logical_to_all_physical_map[ + layer_id, : + ], + partial_logical_to_all_physical_map_num_valid=expert_location_metadata.logical_to_all_physical_map_num_valid[ + layer_id, : + ], + num_physical_experts=expert_location_metadata.num_physical_experts, + ) + + +def transform_select_experts_inputs( + router_logits: torch.Tensor, + correction_bias: Optional[torch.Tensor], + info: Optional[ExpertLocationDispatchInfo], +): + if (info is not None) and (info.ep_dispatch_algorithm == "fake"): + router_logits.uniform_(5, 10) + if correction_bias is not None: + correction_bias = torch.zeros_like(correction_bias) + return router_logits, correction_bias + + +def topk_ids_logical_to_physical( + topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo] +) -> torch.Tensor: + if info is None: + return topk_ids + + if info.ep_dispatch_algorithm == "static": + return _topk_ids_logical_to_physical_static(topk_ids, info) + if info.ep_dispatch_algorithm in ["dynamic", "fake"]: + return _topk_ids_logical_to_physical_dynamic(topk_ids, info) + raise NotImplementedError(f"Unknown algorithm {info.ep_dispatch_algorithm}") + + +def _topk_ids_logical_to_physical_static( + topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo] +) -> torch.Tensor: + return info.partial_logical_to_rank_dispatch_physical_map[topk_ids] + + +def _topk_ids_logical_to_physical_dynamic( + topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo] +) -> torch.Tensor: + topk_ids_original_shape = topk_ids.shape + device = topk_ids.device + topk_ids = topk_ids.flatten() + + chosen_dispatch_index = ( + torch.randint(0, 65536, topk_ids.shape, dtype=torch.int32, device=device) + % info.partial_logical_to_all_physical_map_num_valid[topk_ids] + ) + topk_ids = info.partial_logical_to_all_physical_map[topk_ids, chosen_dispatch_index] + + topk_ids = topk_ids.view(topk_ids_original_shape) + return topk_ids diff --git a/python/sglang/srt/eplb/expert_location_updater.py b/python/sglang/srt/eplb/expert_location_updater.py new file mode 100644 index 000000000..772e65f18 --- /dev/null +++ b/python/sglang/srt/eplb/expert_location_updater.py @@ -0,0 +1,575 @@ +# Copyright 2023-2025 SGLang 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 logging +from collections import defaultdict +from typing import Dict, List, Optional, Tuple + +import einops +import torch +import torch.distributed +from torch.distributed import P2POp + +from sglang.srt.eplb.expert_location import ( + ExpertLocationMetadata, + get_global_expert_location_metadata, +) +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.utils import get_bool_env_var + +logger = logging.getLogger(__name__) + + +_LOG_INPUT = get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_LOG_INPUT") + + +class ExpertLocationUpdater: + def __init__(self): + self._first_execution = True + + def update( + self, + routed_experts_weights_of_layer: Dict[int, List[torch.Tensor]], + new_expert_location_metadata: ExpertLocationMetadata, + update_layer_ids: List[int], + nnodes: int, + rank: int, + ): + if self._first_execution: + self._first_execution = False + torch.get_device_module().empty_cache() + + old_expert_location_metadata = get_global_expert_location_metadata() + assert old_expert_location_metadata is not None + + _update_expert_weights( + routed_experts_weights_of_layer=routed_experts_weights_of_layer, + old_expert_location_metadata=old_expert_location_metadata, + new_expert_location_metadata=new_expert_location_metadata, + update_layer_ids=update_layer_ids, + nnodes=nnodes, + rank=rank, + ) + old_expert_location_metadata.update( + new_expert_location_metadata, + update_layer_ids=update_layer_ids, + ) + + +def _update_expert_weights(**kwargs): + if get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_CANARY"): + return _update_expert_weights_with_canary(**kwargs) + else: + return _update_expert_weights_raw(**kwargs) + + +# can add watchdog as well +def _update_expert_weights_with_canary( + routed_experts_weights_of_layer: Dict[int, List[torch.Tensor]], + old_expert_location_metadata: ExpertLocationMetadata, + new_expert_location_metadata: ExpertLocationMetadata, + update_layer_ids: List[int], + nnodes: int, + rank: int, +): + num_local_physical_experts = old_expert_location_metadata.num_local_physical_experts + + def _get_canary_value(meta: ExpertLocationMetadata, layer_id: int): + return meta.physical_to_logical_map_cpu[ + layer_id, + num_local_physical_experts * rank : num_local_physical_experts * (rank + 1), + ] + + routed_experts_weights_of_layer = { + k: [x for x in v] for k, v in routed_experts_weights_of_layer.items() + } + for layer_id in update_layer_ids: + canary_tensor = ( + _get_canary_value(old_expert_location_metadata, layer_id) + .clone() + .to(device=global_server_args_dict["device"], non_blocking=True) + ) + routed_experts_weights_of_layer[layer_id].append(canary_tensor) + + _update_expert_weights_raw( + routed_experts_weights_of_layer=routed_experts_weights_of_layer, + old_expert_location_metadata=old_expert_location_metadata, + new_expert_location_metadata=new_expert_location_metadata, + update_layer_ids=update_layer_ids, + nnodes=nnodes, + rank=rank, + ) + + for layer_id in update_layer_ids: + # can optimize speed if needed + expect_value = _get_canary_value(new_expert_location_metadata, layer_id) + actual_value = routed_experts_weights_of_layer[layer_id][-1].cpu() + assert torch.all(expect_value == actual_value), ( + f"{expect_value=} {actual_value=} {layer_id=} " + f"{old_expert_location_metadata.physical_to_logical_map_cpu.tolist()=} " + f"{new_expert_location_metadata.physical_to_logical_map_cpu.tolist()=} " + ) + + +def _update_expert_weights_raw( + routed_experts_weights_of_layer: Dict[int, List[torch.Tensor]], + old_expert_location_metadata: ExpertLocationMetadata, + new_expert_location_metadata: ExpertLocationMetadata, + update_layer_ids: List[int], + nnodes: int, + rank: int, +): + log_metrics = get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_LOG_METRICS") + + temp_buffers = create_temp_buffers( + routed_experts_weights_of_layer[update_layer_ids[0]] + ) + + world_size = torch.distributed.get_world_size() + num_local_physical_experts = old_expert_location_metadata.num_local_physical_experts + num_gpu_per_node = world_size // nnodes + + for layer_id in update_layer_ids: + update_expert_weights_single_layer( + routed_experts_weights=routed_experts_weights_of_layer[layer_id], + temp_buffers=temp_buffers, + old_physical_to_logical_map=old_expert_location_metadata.physical_to_logical_map_cpu[ + layer_id + ].tolist(), + new_physical_to_logical_map=new_expert_location_metadata.physical_to_logical_map_cpu[ + layer_id + ].tolist(), + num_local_physical_experts=num_local_physical_experts, + num_gpu_per_node=num_gpu_per_node, + rank=rank, + world_size=world_size, + log_metrics=log_metrics, + ) + + +def create_temp_buffers(sample_tensors): + return [torch.empty_like(tensor) for tensor in sample_tensors] + + +def update_expert_weights_single_layer( + routed_experts_weights: List[torch.Tensor], + temp_buffers: List[torch.Tensor], + old_physical_to_logical_map: List[int], # (num_physical_Experts,) + new_physical_to_logical_map: List[int], # (num_physical_Experts,) + num_local_physical_experts: int, + num_gpu_per_node: int, + rank: int, + world_size: Optional[int] = None, + debug: bool = False, + log_metrics: bool = False, +): + assert all( + tensor.shape[0] == num_local_physical_experts + for tensor in routed_experts_weights + ), f"{num_local_physical_experts=} {[x.shape for x in routed_experts_weights]=}" + assert isinstance(old_physical_to_logical_map, list) + assert isinstance(new_physical_to_logical_map, list) + + if _LOG_INPUT: + logger.info( + "update_expert_weights_single_layer " + f"{[x.shape for x in routed_experts_weights]=} " + f"{[x.shape for x in temp_buffers]=} " + f"{old_physical_to_logical_map=} " + f"{new_physical_to_logical_map=} " + f"{num_local_physical_experts=} " + f"{num_gpu_per_node=} " + f"{rank=} " + f"{world_size=} " + ) + + output_logs = [] if debug else None + + num_physical_experts = len(old_physical_to_logical_map) + num_tensors = len(routed_experts_weights) + + self_node_id = rank // num_gpu_per_node + + local_expert_location_range = ( + rank * num_local_physical_experts, + (rank + 1) * num_local_physical_experts, + ) + + def _entrypoint(): + # List[Tuple[logical_expert_id, List[P2POp]]] + p2p_op_infos: List[Tuple[int, List[P2POp]]] = [] + # List[Tuple[temp_buffers_expert_location, routed_experts_weights_expert_location]] + buffer2weight_copy_infos: List[Tuple[int, int]] = [] + + _handle_recv(buffer2weight_copy_infos, p2p_op_infos) + _create_isend_ops(p2p_op_infos) + _execute_p2p_ops(p2p_op_infos) + _execute_buffer2weight_copies(buffer2weight_copy_infos) + + if log_metrics: + _log_p2p_op_metrics( + p2p_op_infos, + world_size=world_size, + num_gpu_per_node=num_gpu_per_node, + self_node_id=self_node_id, + ) + + if debug: + output_logs.append(f"{p2p_op_infos=}") + output_logs.append(f"{buffer2weight_copy_infos=}") + + def _handle_recv(buffer2weight_copy_infos, p2p_op_infos): + for dst_expert_location in range(*local_expert_location_range): + _handle_recv_of_dst_expert_location( + dst_expert_location, buffer2weight_copy_infos, p2p_op_infos + ) + + def _handle_recv_of_dst_expert_location( + dst_expert_location: int, buffer2weight_copy_infos, p2p_op_infos + ): + logical_expert_id = new_physical_to_logical_map[dst_expert_location] + + # case 1: unchanged + if old_physical_to_logical_map[dst_expert_location] == logical_expert_id: + if debug: + output_logs.append( + f"handle_recv_of_dst_expert_location {dst_expert_location=} case=unchanged" + ) + return + + # case 2: same-gpu + for src_expert_location in range(*local_expert_location_range): + if old_physical_to_logical_map[src_expert_location] == logical_expert_id: + for i in range(num_tensors): + _get_tensor(temp_buffers, i, dst_expert_location).copy_( + _get_tensor(routed_experts_weights, i, src_expert_location) + ) + buffer2weight_copy_infos.append( + (dst_expert_location, dst_expert_location) + ) + if debug: + output_logs.append( + f"handle_recv_of_dst_expert_location {dst_expert_location=} case=same-gpu {src_expert_location=}" + ) + return + + # case 3: free-rider + for src_expert_location in range( + rank * num_local_physical_experts, dst_expert_location + ): + if new_physical_to_logical_map[src_expert_location] == logical_expert_id: + buffer2weight_copy_infos.append( + (src_expert_location, dst_expert_location) + ) + if debug: + output_logs.append( + f"handle_recv_of_dst_expert_location {dst_expert_location=} case=free-rider {src_expert_location=}" + ) + return + + same_node_mapping, cross_node_mapping, need_comm_self_node_dst_ranks = ( + _compute_comm_info(logical_expert_id=logical_expert_id) + ) + + # case 4: same-node + if rank in need_comm_self_node_dst_ranks: + chosen_src_rank = same_node_mapping.chunk_value_from_element_value( + element_value=rank + ) + _create_p2p_recv_and_buffer2weight_copy( + buffer2weight_copy_infos, + p2p_op_infos, + src_rank=chosen_src_rank, + logical_expert_id=logical_expert_id, + dst_expert_location=dst_expert_location, + ) + if debug: + output_logs.append( + f"handle_recv_of_dst_expert_location {dst_expert_location=} case=same-node {chosen_src_rank=}" + ) + return + + # case 5: cross-node + # Future work: can optimize when there are multiple ranks in the same dst node that uses the same logical expert + chosen_src_rank = cross_node_mapping.chunk_value_from_element_value( + element_value=rank + ) + _create_p2p_recv_and_buffer2weight_copy( + buffer2weight_copy_infos, + p2p_op_infos, + src_rank=chosen_src_rank, + logical_expert_id=logical_expert_id, + dst_expert_location=dst_expert_location, + ) + if debug: + output_logs.append( + f"handle_recv_of_dst_expert_location {dst_expert_location=} case=cross-node {chosen_src_rank=}" + ) + return + + def _create_p2p_recv_and_buffer2weight_copy( + buffer2weight_copy_infos, + p2p_op_infos, + *, + logical_expert_id: int, + src_rank: int, + dst_expert_location: int, + ): + p2p_op_infos.append( + ( + logical_expert_id, + [ + P2POp( + op=torch.distributed.irecv, + tensor=_get_tensor(temp_buffers, i, dst_expert_location), + peer=src_rank, + ) + for i in range(num_tensors) + ], + ) + ) + buffer2weight_copy_infos.append((dst_expert_location, dst_expert_location)) + + def _create_isend_ops(p2p_op_infos): + handled_logical_expert_ids = set() + for src_expert_location in range(*local_expert_location_range): + logical_expert_id = old_physical_to_logical_map[src_expert_location] + + if logical_expert_id in handled_logical_expert_ids: + continue + handled_logical_expert_ids.add(logical_expert_id) + + _create_isend_ops_of_logical_expert_id( + logical_expert_id, src_expert_location, p2p_op_infos + ) + + def _create_isend_ops_of_logical_expert_id( + logical_expert_id, src_expert_location, p2p_op_infos + ): + same_node_mapping, cross_node_mapping, need_comm_self_node_dst_ranks = ( + _compute_comm_info(logical_expert_id=logical_expert_id) + ) + + same_node_dst_ranks = same_node_mapping.element_values_from_chunk_value( + chunk_value=rank + ) + cross_node_dst_ranks = cross_node_mapping.element_values_from_chunk_value( + chunk_value=rank + ) + all_dst_ranks = same_node_dst_ranks + cross_node_dst_ranks + + if debug: + output_logs.append( + f"create_isend_ops_of_logical_expert_id {logical_expert_id=} {src_expert_location=} {same_node_dst_ranks=} {cross_node_dst_ranks=}" + ) + + p2p_op_infos.append( + ( + logical_expert_id, + [ + P2POp( + op=torch.distributed.isend, + tensor=_get_tensor( + routed_experts_weights, i, src_expert_location + ), + peer=dst_rank, + ) + for dst_rank in all_dst_ranks + for i in range(num_tensors) + ], + ) + ) + + def _compute_comm_info(logical_expert_id: int): + all_src_ranks = _deduplicate_ordered( + [ + x // num_local_physical_experts + for x in range(num_physical_experts) + if old_physical_to_logical_map[x] == logical_expert_id + ] + ) + all_src_nodes = [x // num_gpu_per_node for x in all_src_ranks] + self_node_src_ranks = [ + x for x in all_src_ranks if x // num_gpu_per_node == self_node_id + ] + + need_comm_dst_ranks = _deduplicate_ordered( + [ + x // num_local_physical_experts + for x in range(num_physical_experts) + if new_physical_to_logical_map[x] == logical_expert_id + and x // num_local_physical_experts not in all_src_ranks + ] + ) + need_comm_self_node_dst_ranks = ( + [x for x in need_comm_dst_ranks if x // num_gpu_per_node == self_node_id] + if len(self_node_src_ranks) > 0 + else [] + ) + need_comm_cross_node_dst_ranks = [ + x + for x in need_comm_dst_ranks + if (x // num_gpu_per_node) not in all_src_nodes + ] + + same_node_mapping = _ChunkUtils( + chunk_values=self_node_src_ranks, + element_values=need_comm_self_node_dst_ranks, + ) + + cross_node_mapping = _ChunkUtils( + chunk_values=all_src_ranks, + element_values=need_comm_cross_node_dst_ranks, + ) + + return same_node_mapping, cross_node_mapping, need_comm_self_node_dst_ranks + + def _execute_p2p_ops(p2p_op_infos): + sorted_infos = sorted(p2p_op_infos, key=lambda info: info[0]) + p2p_ops = [op for _, ops in sorted_infos for op in ops] + if len(p2p_ops) == 0: + return + + reqs = torch.distributed.batch_isend_irecv(p2p_ops) + for req in reqs: + req.wait() + + def _execute_buffer2weight_copies(buffer2weight_copy_infos): + for ( + temp_buffers_expert_location, + routed_experts_weights_expert_location, + ) in buffer2weight_copy_infos: + for i in range(num_tensors): + _get_tensor( + routed_experts_weights, i, routed_experts_weights_expert_location + ).copy_(_get_tensor(temp_buffers, i, temp_buffers_expert_location)) + + def _get_tensor(tensors, tensor_index: int, expert_location: int) -> torch.Tensor: + return tensors[tensor_index][_get_local_expert_location(expert_location)] + + def _get_local_expert_location(expert_location: int) -> int: + assert ( + local_expert_location_range[0] + <= expert_location + < local_expert_location_range[1] + ) + return expert_location % num_local_physical_experts + + _entrypoint() + + return output_logs + + +class _ChunkUtils: + def __init__(self, *, chunk_values: List, element_values: List): + self.chunk_values = chunk_values + self.element_values = element_values + + def chunk_value_from_element_value(self, element_value): + chunk_index = self._chunk_index_from_element_index( + num_elements=len(self.element_values), + num_chunks=len(self.chunk_values), + element_index=self.element_values.index(element_value), + ) + return self.chunk_values[chunk_index] + + def element_values_from_chunk_value(self, chunk_value) -> List: + if len(self.element_values) == 0: + return [] + element_slice = self._element_slice_from_chunk_index( + num_elements=len(self.element_values), + num_chunks=len(self.chunk_values), + chunk_index=self.chunk_values.index(chunk_value), + ) + return self.element_values[element_slice] + + @staticmethod + def _chunk_index_from_element_index( + num_elements: int, num_chunks: int, element_index: int + ) -> int: + short_chunk_size, num_long_chunks = divmod(num_elements, num_chunks) + num_elements_for_long_chunks = num_long_chunks * (short_chunk_size + 1) + if element_index < num_elements_for_long_chunks: + return element_index // (short_chunk_size + 1) + else: + return ( + num_long_chunks + + (element_index - num_elements_for_long_chunks) // short_chunk_size + ) + + @staticmethod + def _element_slice_from_chunk_index( + num_elements: int, num_chunks: int, chunk_index: int + ) -> slice: + short_chunk_size, num_long_chunks = divmod(num_elements, num_chunks) + start = chunk_index * short_chunk_size + min(chunk_index, num_long_chunks) + end = start + short_chunk_size + int(chunk_index < num_long_chunks) + return slice(start, end) + + +def _deduplicate_ordered(arr: List[int]): + output = [] + for item in arr: + if len(output) == 0 or item != output[-1]: + output.append(item) + return output + + +def _log_p2p_op_metrics( + p2p_op_infos: List[Tuple[int, List[P2POp]]], + num_gpu_per_node: int, + world_size: int, + self_node_id: int, +): + text = "" + all_ops = [op for _, ops in p2p_op_infos for op in ops] + + for direction, ops in _group_by(all_ops, _get_direction_from_op).items(): + nbytes_of_gpu = [0] * world_size + for op in ops: + nbytes_of_gpu[op.peer] += op.tensor.nbytes + nbytes_of_gpu = torch.tensor(nbytes_of_gpu, dtype=torch.int64) + + nbytes_of_node = einops.reduce( + nbytes_of_gpu, + "(num_nodes num_gpu_per_node) -> num_nodes", + num_gpu_per_node=num_gpu_per_node, + reduction="sum", + ) + + nbytes_curr_node = nbytes_of_node[self_node_id] + nbytes_cross_node = torch.sum(nbytes_of_node) - nbytes_curr_node + + text += ( + f"{direction}_nbytes_of_gpu={nbytes_of_gpu.tolist()} " + f"{direction}_nbytes_of_node={nbytes_of_node.tolist()} " + f"{direction}_nbytes_curr_node={nbytes_curr_node.item()} " + f"{direction}_nbytes_cross_node={nbytes_cross_node.item()} " + ) + + logger.info(f"[ExpertLocationUpdater] {text}") + + +def _get_direction_from_op(op: P2POp): + if op.op == torch.distributed.isend: + return "isend" + if op.op == torch.distributed.irecv: + return "irecv" + raise NotImplementedError + + +def _group_by(items, keyfunc): + ans = defaultdict(list) + for item in items: + ans[keyfunc(item)].append(item) + return dict(ans) diff --git a/python/sglang/srt/function_call/base_format_detector.py b/python/sglang/srt/function_call/base_format_detector.py new file mode 100644 index 000000000..39bb92f5f --- /dev/null +++ b/python/sglang/srt/function_call/base_format_detector.py @@ -0,0 +1,367 @@ +import json +import logging +from abc import ABC, abstractmethod +from typing import Any, Dict, List + +from partial_json_parser.core.exceptions import MalformedJSON +from partial_json_parser.core.options import Allow + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.function_call.utils import ( + _find_common_prefix, + _is_complete_json, + _partial_json_loads, +) + +logger = logging.getLogger(__name__) + + +class BaseFormatDetector(ABC): + """Base class providing two sets of interfaces: one-time and streaming incremental.""" + + def __init__(self): + # Streaming state management + # Buffer for accumulating incomplete patterns that arrive across multiple streaming chunks + self._buffer = "" + # Stores complete tool call info (name and arguments) for each tool being parsed. + # Used by serving layer for completion handling when streaming ends. + # Format: [{"name": str, "arguments": dict}, ...] + self.prev_tool_call_arr: List[Dict] = [] + # Index of currently streaming tool call. Starts at -1 (no active tool), + # increments as each tool completes. Tracks which tool's arguments are streaming. + self.current_tool_id: int = -1 + # Flag for whether current tool's name has been sent to client. + # Tool names sent first with empty parameters, then arguments stream incrementally. + self.current_tool_name_sent: bool = False + # Tracks raw JSON string content streamed to client for each tool's arguments. + # Critical for serving layer to calculate remaining content when streaming ends. + # Each index corresponds to a tool_id. Example: ['{"location": "San Francisco"', '{"temp": 72'] + self.streamed_args_for_tool: List[str] = [] + + # Token configuration (override in subclasses) + self.bot_token = "" + self.eot_token = "" + self.tool_call_separator = ", " + + def _get_tool_indices(self, tools: List[Tool]) -> Dict[str, int]: + """ + Get a mapping of tool names to their indices in the tools list. + + This utility method creates a dictionary mapping function names to their + indices in the tools list, which is commonly needed for tool validation + and ToolCallItem creation. + + Args: + tools: List of available tools + + Returns: + Dictionary mapping tool names to their indices + """ + return { + tool.function.name: i for i, tool in enumerate(tools) if tool.function.name + } + + def parse_base_json(self, action: Any, tools: List[Tool]) -> List[ToolCallItem]: + tool_indices = self._get_tool_indices(tools) + if not isinstance(action, list): + action = [action] + + results = [] + for act in action: + name = act.get("name") + if name and name in tool_indices: + results.append( + ToolCallItem( + tool_index=-1, # Caller should update this based on the actual tools array called + name=name, + parameters=json.dumps( + act.get("parameters") or act.get("arguments", {}), + ensure_ascii=False, + ), + ) + ) + else: + logger.warning(f"Model attempted to call undefined function: {name}") + + return results + + @abstractmethod + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """ + Parses the text in one go. Returns success=True if the format matches, otherwise False. + Note that leftover_text here represents "content that this parser will not consume further". + """ + action = json.loads(text) + return StreamingParseResult(calls=self.parse_base_json(action, tools)) + + def _ends_with_partial_token(self, buffer: str, bot_token: str) -> int: + """ + Check if buffer ends with a partial bot_token. + Return the length of the partial bot_token. + + For some format, the bot_token is not a token in model's vocabulary, such as + `[TOOL_CALLS] [` in Mistral. + """ + for i in range(1, min(len(buffer) + 1, len(bot_token))): + if bot_token.startswith(buffer[-i:]): + return i + return 0 + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Streaming incremental parsing with tool validation. + + This base implementation works best with formats where: + 1. bot_token is followed immediately by JSON (e.g., bot_token + JSON_array) + 2. JSON can be parsed incrementally using partial_json_loads + 3. Multiple tool calls are separated by "; " or ", " + + Examples of incompatible formats (need custom implementation, may reuse some logic from this class): + - Each tool call is wrapped in a separate block: See Qwen25Detector + - Multiple separate blocks: [TOOL_CALLS] [...] \n [TOOL_CALLS] [...] + - Tool call is Pythonic style + + For incompatible formats, detectors should override this method with custom logic. + """ + # Append new text to buffer + self._buffer += new_text + current_text = self._buffer + + # The current_text has tool_call if it is the start of a new tool call sequence + # or it is the start of a new tool call after a tool call separator, when there is a previous tool call + if not ( + self.has_tool_call(current_text) + or ( + self.current_tool_id > 0 + and current_text.startswith(self.tool_call_separator) + ) + ): + # Only clear buffer if we're sure no tool call is starting + if not self._ends_with_partial_token(self._buffer, self.bot_token): + normal_text = self._buffer + self._buffer = "" + if self.eot_token in normal_text: + normal_text = normal_text.replace(self.eot_token, "") + return StreamingParseResult(normal_text=normal_text) + else: + # Might be partial bot_token, keep buffering + return StreamingParseResult() + + # Build tool indices if not already built + if not hasattr(self, "_tool_indices"): + self._tool_indices = self._get_tool_indices(tools) + + flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR + + try: + try: + if current_text.startswith(self.bot_token): + start_idx = len(self.bot_token) + elif self.current_tool_id > 0 and current_text.startswith( + self.tool_call_separator + self.bot_token + ): + start_idx = len(self.tool_call_separator + self.bot_token) + elif self.current_tool_id > 0 and current_text.startswith( + self.tool_call_separator + ): + start_idx = len(self.tool_call_separator) + else: + start_idx = 0 + + if start_idx >= len(current_text): + return StreamingParseResult() + + (obj, end_idx) = _partial_json_loads(current_text[start_idx:], flags) + + is_current_complete = _is_complete_json( + current_text[start_idx : start_idx + end_idx] + ) + + # Validate tool name if present + if "name" in obj and obj["name"] not in self._tool_indices: + # Invalid tool name - reset state + self._buffer = "" + self.current_tool_id = -1 + self.current_tool_name_sent = False + if self.streamed_args_for_tool: + self.streamed_args_for_tool.pop() + return StreamingParseResult() + + # Handle parameters/arguments consistency + # NOTE: we assume here that the obj is always partial of a single tool call + if "parameters" in obj: + assert ( + "arguments" not in obj + ), "model generated both parameters and arguments" + obj["arguments"] = obj["parameters"] + + current_tool_call = obj + + except MalformedJSON: + return StreamingParseResult() + + if not current_tool_call: + return StreamingParseResult() + + # Case 1: Handle tool name streaming + # This happens when we encounter a tool but haven't sent its name yet + if not self.current_tool_name_sent: + function_name = current_tool_call.get("name") + + if function_name and function_name in self._tool_indices: + # If this is a new tool (current_tool_id was -1), initialize it + if self.current_tool_id == -1: + self.current_tool_id = 0 + self.streamed_args_for_tool.append("") + # If this is a subsequent tool, ensure streamed_args_for_tool is large enough + elif self.current_tool_id >= len(self.streamed_args_for_tool): + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + + # Send the tool name with empty parameters + res = StreamingParseResult( + calls=[ + ToolCallItem( + tool_index=self.current_tool_id, + name=function_name, + parameters="", + ) + ], + ) + self.current_tool_name_sent = True + else: + res = StreamingParseResult() + + # Case 2: Handle streaming arguments + # This happens when we've already sent the tool name and now need to stream arguments incrementally + else: + cur_arguments = current_tool_call.get("arguments") + res = StreamingParseResult() + + if cur_arguments: + # Calculate how much of the arguments we've already streamed + sent = len(self.streamed_args_for_tool[self.current_tool_id]) + cur_args_json = json.dumps(cur_arguments) + prev_arguments = None + if self.current_tool_id < len(self.prev_tool_call_arr): + prev_arguments = self.prev_tool_call_arr[ + self.current_tool_id + ].get("arguments") + + argument_diff = None + + # If the current tool's JSON is complete, send all remaining arguments + if is_current_complete: + argument_diff = cur_args_json[sent:] + completing_tool_id = ( + self.current_tool_id + ) # Save the ID of the tool that's completing + + # Only remove the processed portion, keep unprocessed content + self._buffer = current_text[start_idx + end_idx :] + + if self.current_tool_id < len(self.prev_tool_call_arr): + self.prev_tool_call_arr[self.current_tool_id].clear() + self.current_tool_name_sent = False + self.streamed_args_for_tool[self.current_tool_id] = "" + self.current_tool_id += 1 + + # If the tool is still being parsed, send incremental changes + elif prev_arguments: + prev_args_json = json.dumps(prev_arguments) + if cur_args_json != prev_args_json: + prefix = _find_common_prefix(prev_args_json, cur_args_json) + argument_diff = prefix[sent:] + + # Send the argument diff if there's something new + if argument_diff is not None: + # Use the correct tool_index: completing_tool_id for completed tools, current_tool_id for ongoing + tool_index_to_use = ( + completing_tool_id + if is_current_complete + else self.current_tool_id + ) + res = StreamingParseResult( + calls=[ + ToolCallItem( + tool_index=tool_index_to_use, + parameters=argument_diff, + ) + ], + ) + if not is_current_complete: + self.streamed_args_for_tool[ + self.current_tool_id + ] += argument_diff + + # Update prev_tool_call_arr with current state + if self.current_tool_id >= 0: + # Ensure prev_tool_call_arr is large enough + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + self.prev_tool_call_arr[self.current_tool_id] = current_tool_call + + return res + + except Exception as e: + logger.error(f"Error in parse_streaming_increment: {e}") + return StreamingParseResult() + + @abstractmethod + def has_tool_call(self, text: str) -> bool: + """ + Check if the given text contains function call markers specific to this format. + """ + raise NotImplementedError() + + def supports_structural_tag(self) -> bool: + """Return True if this detector supports structural tag format.""" + return True + + @abstractmethod + def structure_info(self) -> _GetInfoFunc: + """ + Return a function that creates StructureInfo for constrained generation. + + The returned function takes a tool name and returns a StructureInfo object + containing the begin/end patterns and trigger tokens needed for constrained + generation of function calls in this format. + + Returns: + A function that takes a tool name (str) and returns StructureInfo + """ + raise NotImplementedError() + + @abstractmethod + def build_ebnf(self, tools: List[Tool]) -> str: + """ + Build an EBNF grammar for constrained generation of function calls. + + This method generates an Extended Backus-Naur Form (EBNF) grammar that + constrains the model's output to valid function calls in this format. + The grammar should include all available tools and their parameter schemas. + + Args: + tools: List of available tools/functions that can be called + + Returns: + A string containing the EBNF grammar for this function call format + + The EBNF grammar should: + - Define the overall structure of function calls in this format + - Include all tool names from the provided tools list + - Define valid JSON structures for function arguments + - Handle multiple function calls if the format supports them + + Note: + Most implementations use EBNFComposer.build_ebnf() utility with + format-specific parameters rather than writing EBNF from scratch. + """ + raise NotImplementedError() diff --git a/python/sglang/srt/function_call/core_types.py b/python/sglang/srt/function_call/core_types.py new file mode 100644 index 000000000..1ea87df79 --- /dev/null +++ b/python/sglang/srt/function_call/core_types.py @@ -0,0 +1,34 @@ +from dataclasses import dataclass +from typing import Callable, List, Optional + +from pydantic import BaseModel + + +class ToolCallItem(BaseModel): + """Simple encapsulation of the parsed ToolCall result for easier usage in streaming contexts.""" + + tool_index: int + name: Optional[str] = None + parameters: str # JSON string + + +class StreamingParseResult(BaseModel): + """Result of streaming incremental parsing.""" + + normal_text: str = "" + calls: List[ToolCallItem] = [] + + +@dataclass +class StructureInfo: + begin: str + end: str + trigger: str + + +""" +Helper alias of function +Usually it is a function that takes a name string and returns a StructureInfo object, +which can be used to construct a structural_tag object +""" +_GetInfoFunc = Callable[[str], StructureInfo] diff --git a/python/sglang/srt/function_call/deepseekv31_detector.py b/python/sglang/srt/function_call/deepseekv31_detector.py new file mode 100644 index 000000000..2045d8daa --- /dev/null +++ b/python/sglang/srt/function_call/deepseekv31_detector.py @@ -0,0 +1,222 @@ +import json +import logging +import re +from typing import List + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer +from sglang.srt.function_call.utils import _is_complete_json + +logger = logging.getLogger(__name__) + + +class DeepSeekV31Detector(BaseFormatDetector): + """ + Detector for DeepSeek V3 model function call format. + + The DeepSeek V3 format uses special Unicode tokens to delimit function calls + with JSON code blocks for arguments. + + Format Structure: + ``` + <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>{function_name}<｜tool▁sep｜>{json_arguments}<｜tool▁calls▁end｜><｜end▁of▁sentence｜> + ``` + Examples: + ``` + <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_current_weather<｜tool▁sep｜>{"location": "Tokyo"}<｜tool▁call▁end｜><｜tool▁call▁begin｜>get_current_weather<｜tool▁sep｜>{"location": "Paris"}<｜tool▁call▁end｜><｜tool▁calls▁end｜><｜end▁of▁sentence｜> + ``` + + Key Components: + - Tool Calls Section: Wrapped between `<｜tool▁calls▁begin｜>` and `<｜tool▁calls▁end｜>` + - Individual Tool Call: Wrapped between `<｜tool▁call▁begin｜>` and `<｜tool▁call▁end｜>` + - Function Declaration: `<｜tool▁call▁begin｜>{function_name}<｜tool▁sep｜>` + - Arguments: JSON code block between `<｜tool▁sep｜>` and `<｜tool▁call▁end｜>` + - Supports multiple tool calls + + Reference: https://www.modelscope.cn/models/deepseek-ai/DeepSeek-V3.1 + """ + + def __init__(self): + super().__init__() + self.bot_token = "<｜tool▁calls▁begin｜>" + self.eot_token = "<｜tool▁calls▁end｜>" + self.func_call_regex = r"<｜tool▁call▁begin｜>.*?<｜tool▁call▁end｜>" + self.func_detail_regex = ( + r"<｜tool▁call▁begin｜>(.*)<｜tool▁sep｜>(.*)<｜tool▁call▁end｜>" + ) + self._last_arguments = "" + self.current_tool_id = -1 + + def has_tool_call(self, text: str) -> bool: + """Check if the text contains a deepseek format tool call.""" + return self.bot_token in text + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """ + One-time parsing: Detects and parses tool calls in the provided text. + + :param text: The complete text to parse. + :param tools: List of available tools. + :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. + """ + idx = text.find(self.bot_token) + normal_text = text[:idx].strip() if idx != -1 else text + if self.bot_token not in text: + return StreamingParseResult(normal_text=normal_text, calls=[]) + match_result_list = re.findall(self.func_call_regex, text, re.DOTALL) + calls = [] + try: + for match_result in match_result_list: + # Get function name + func_detail = re.search(self.func_detail_regex, match_result, re.DOTALL) + func_name = func_detail.group(1) + func_args = func_detail.group(2) + func_args = json.loads(func_args) + # construct match_result for parse_base_json + match_result = {"name": func_name, "parameters": func_args} + calls.extend(self.parse_base_json(match_result, tools)) + return StreamingParseResult(normal_text=normal_text, calls=calls) + except Exception as e: + logger.error(f"Error in detect_and_parse: {e}") + # return the normal text if parsing fails + return StreamingParseResult(normal_text=text) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Streaming incremental parsing tool calls for DeepSeekV3 format. + """ + self._buffer += new_text + current_text = self._buffer + + # Check if we have a tool call (either the start token or individual tool call) + has_tool_call = ( + self.bot_token in current_text or "<｜tool▁call▁begin｜>" in current_text + ) + + if not has_tool_call: + self._buffer = "" + for e_token in [self.eot_token, "<｜tool▁call▁end｜>"]: + if e_token in new_text: + new_text = new_text.replace(e_token, "") + return StreamingParseResult(normal_text=new_text) + + if not hasattr(self, "_tool_indices"): + self._tool_indices = self._get_tool_indices(tools) + + calls: list[ToolCallItem] = [] + try: + partial_match = re.search( + pattern=r"<｜tool▁call▁begin｜>(.*)<｜tool▁sep｜>(.*)<｜tool▁call▁end｜>", + string=current_text, + flags=re.DOTALL, + ) + if partial_match: + func_name = partial_match.group(1).strip() + func_args_raw = partial_match.group(2).strip() + + # Initialize state if this is the first tool call + if self.current_tool_id == -1: + self.current_tool_id = 0 + self.prev_tool_call_arr = [] + self.streamed_args_for_tool = [""] + + # Ensure we have enough entries in our tracking arrays + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + + if not self.current_tool_name_sent: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=func_name, + parameters="", + ) + ) + self.current_tool_name_sent = True + # Store the tool call info for serving layer completions endpoint + self.prev_tool_call_arr[self.current_tool_id] = { + "name": func_name, + "arguments": {}, + } + else: + argument_diff = ( + func_args_raw[len(self._last_arguments) :] + if func_args_raw.startswith(self._last_arguments) + else func_args_raw + ) + + if argument_diff: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=None, + parameters=argument_diff, + ) + ) + self._last_arguments += argument_diff + self.streamed_args_for_tool[ + self.current_tool_id + ] += argument_diff + + if _is_complete_json(func_args_raw): + # Update the stored arguments + try: + parsed_args = json.loads(func_args_raw) + self.prev_tool_call_arr[self.current_tool_id][ + "arguments" + ] = parsed_args + except json.JSONDecodeError: + pass + + # Find the end of the current tool call and remove only that part from buffer + tool_call_end_pattern = ( + r"<｜tool▁call▁begin｜>.*?<｜tool▁call▁end｜>" + ) + match = re.search( + tool_call_end_pattern, current_text, re.DOTALL + ) + if match: + # Remove the completed tool call from buffer, keep any remaining content + self._buffer = current_text[match.end() :] + else: + self._buffer = "" + + result = StreamingParseResult(normal_text="", calls=calls) + self.current_tool_id += 1 + self._last_arguments = "" + self.current_tool_name_sent = False + return result + + return StreamingParseResult(normal_text="", calls=calls) + + except Exception as e: + logger.error(f"Error in parse_streaming_increment: {e}") + return StreamingParseResult(normal_text=current_text) + + def structure_info(self) -> _GetInfoFunc: + return lambda name: StructureInfo( + begin="<｜tool▁call▁begin｜>" + name + "<｜tool▁sep｜>", + end="<｜tool▁call▁end｜>", + trigger="<｜tool▁call▁begin｜>" + name + "<｜tool▁sep｜>", + ) + + def build_ebnf(self, tools: List[Tool]): + return EBNFComposer.build_ebnf( + tools, + sequence_start_token=self.bot_token, + sequence_end_token=self.eot_token, + tool_call_separator="", + call_rule_fmt='"<｜tool▁call▁begin｜>{name}<｜tool▁sep｜>{arguments_rule}<｜tool▁call▁end｜>"', + function_format="json", + ) diff --git a/python/sglang/srt/function_call/deepseekv3_detector.py b/python/sglang/srt/function_call/deepseekv3_detector.py new file mode 100644 index 000000000..33c4dfc44 --- /dev/null +++ b/python/sglang/srt/function_call/deepseekv3_detector.py @@ -0,0 +1,220 @@ +import json +import logging +import re +from typing import List + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer +from sglang.srt.function_call.utils import _is_complete_json + +logger = logging.getLogger(__name__) + + +class DeepSeekV3Detector(BaseFormatDetector): + """ + Detector for DeepSeek V3 model function call format. + + The DeepSeek V3 format uses special Unicode tokens to delimit function calls + with JSON code blocks for arguments. + + Format Structure: + ``` + <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>{function_name}\n```json\n{json_arguments}\n```<｜tool▁calls▁end｜><｜end▁of▁sentence｜> + ``` + Examples: + ``` + <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_current_weather\n```json\n{"location": "Tokyo"}\n```<｜tool▁call▁end｜>\n<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_current_weather\n```json\n{"location": "Paris"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜><｜end▁of▁sentence｜> + ``` + + Key Components: + - Tool Calls Section: Wrapped between `<｜tool▁calls▁begin｜>` and `<｜tool▁calls▁end｜>` + - Individual Tool Call: Wrapped between `<｜tool▁call▁begin｜>` and `<｜tool▁call▁end｜>` + - Function Declaration: `function<｜tool▁sep｜>{function_name}` + - Arguments: JSON code block between ````json` and ```` + - Supports multiple tool calls + + Reference: https://huggingface.co/deepseek-ai/DeepSeek-V3-0324?chat_template=default + """ + + def __init__(self): + super().__init__() + self.bot_token = "<｜tool▁calls▁begin｜>" + self.eot_token = "<｜tool▁calls▁end｜>" + self.func_call_regex = r"<｜tool▁call▁begin｜>.*?<｜tool▁call▁end｜>" + self.func_detail_regex = r"<｜tool▁call▁begin｜>(.*)<｜tool▁sep｜>(.*)\n```json\n(.*)\n```<｜tool▁call▁end｜>" + self._last_arguments = "" + self.current_tool_id = -1 + + def has_tool_call(self, text: str) -> bool: + """Check if the text contains a deepseek format tool call.""" + return self.bot_token in text + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """ + One-time parsing: Detects and parses tool calls in the provided text. + + :param text: The complete text to parse. + :param tools: List of available tools. + :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. + """ + idx = text.find(self.bot_token) + normal_text = text[:idx].strip() if idx != -1 else text + if self.bot_token not in text: + return StreamingParseResult(normal_text=normal_text, calls=[]) + match_result_list = re.findall(self.func_call_regex, text, re.DOTALL) + calls = [] + try: + for match_result in match_result_list: + # Get function name + func_detail = re.search(self.func_detail_regex, match_result, re.DOTALL) + func_name = func_detail.group(2) + func_args = func_detail.group(3) + func_args = json.loads(func_args) + # construct match_result for parse_base_json + match_result = {"name": func_name, "parameters": func_args} + calls.extend(self.parse_base_json(match_result, tools)) + return StreamingParseResult(normal_text=normal_text, calls=calls) + except Exception as e: + logger.error(f"Error in detect_and_parse: {e}") + # return the normal text if parsing fails + return StreamingParseResult(normal_text=text) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Streaming incremental parsing tool calls for DeepSeekV3 format. + """ + self._buffer += new_text + current_text = self._buffer + + # Check if we have a tool call (either the start token or individual tool call) + has_tool_call = ( + self.bot_token in current_text or "<｜tool▁call▁begin｜>" in current_text + ) + + if not has_tool_call: + self._buffer = "" + for e_token in [self.eot_token, "```", "<｜tool▁call▁end｜>"]: + if e_token in new_text: + new_text = new_text.replace(e_token, "") + return StreamingParseResult(normal_text=new_text) + + if not hasattr(self, "_tool_indices"): + self._tool_indices = self._get_tool_indices(tools) + + calls: list[ToolCallItem] = [] + try: + partial_match = re.search( + pattern=r"<｜tool▁call▁begin｜>(.*)<｜tool▁sep｜>(.*)\n```json\n(.*)\n```.*", + string=current_text, + flags=re.DOTALL, + ) + if partial_match: + func_name = partial_match.group(2).strip() + func_args_raw = partial_match.group(3).strip() + + # Initialize state if this is the first tool call + if self.current_tool_id == -1: + self.current_tool_id = 0 + self.prev_tool_call_arr = [] + self.streamed_args_for_tool = [""] + + # Ensure we have enough entries in our tracking arrays + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + + if not self.current_tool_name_sent: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=func_name, + parameters="", + ) + ) + self.current_tool_name_sent = True + # Store the tool call info for serving layer completions endpoint + self.prev_tool_call_arr[self.current_tool_id] = { + "name": func_name, + "arguments": {}, + } + else: + argument_diff = ( + func_args_raw[len(self._last_arguments) :] + if func_args_raw.startswith(self._last_arguments) + else func_args_raw + ) + + if argument_diff: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=None, + parameters=argument_diff, + ) + ) + self._last_arguments += argument_diff + self.streamed_args_for_tool[ + self.current_tool_id + ] += argument_diff + + if _is_complete_json(func_args_raw): + # Update the stored arguments + try: + parsed_args = json.loads(func_args_raw) + self.prev_tool_call_arr[self.current_tool_id][ + "arguments" + ] = parsed_args + except json.JSONDecodeError: + pass + + # Find the end of the current tool call and remove only that part from buffer + tool_call_end_pattern = ( + r"<｜tool▁call▁begin｜>.*?<｜tool▁call▁end｜>" + ) + match = re.search( + tool_call_end_pattern, current_text, re.DOTALL + ) + if match: + # Remove the completed tool call from buffer, keep any remaining content + self._buffer = current_text[match.end() :] + else: + self._buffer = "" + + result = StreamingParseResult(normal_text="", calls=calls) + self.current_tool_id += 1 + self._last_arguments = "" + self.current_tool_name_sent = False + return result + + return StreamingParseResult(normal_text="", calls=calls) + + except Exception as e: + logger.error(f"Error in parse_streaming_increment: {e}") + return StreamingParseResult(normal_text=current_text) + + def structure_info(self) -> _GetInfoFunc: + return lambda name: StructureInfo( + begin=">" + name + "\n```json\n", + end="\n```<", + trigger=">" + name + "\n```json\n", + ) + + def build_ebnf(self, tools: List[Tool]): + return EBNFComposer.build_ebnf( + tools, + sequence_start_token=self.bot_token, + sequence_end_token=self.eot_token, + tool_call_separator="", + call_rule_fmt='"<｜tool▁call▁begin｜>function<｜tool▁sep｜>{name}\\n```json\\n"{arguments_rule}"\\n```<｜tool▁call▁end｜>"', + function_format="json", + ) diff --git a/python/sglang/srt/function_call/ebnf_composer.py b/python/sglang/srt/function_call/ebnf_composer.py new file mode 100644 index 000000000..21b313982 --- /dev/null +++ b/python/sglang/srt/function_call/ebnf_composer.py @@ -0,0 +1,344 @@ +from typing import Any, Dict, Literal, Optional + + +class EBNFComposer: + # Adapted from https://xgrammar.mlc.ai/docs/how_to/ebnf_guided_generation.html#try-out-via-hf-transformers + # Shared primitive grammar rules used across all formats + BASE_PRIMITIVE_GRAMMAR = r""" + basic_string ::= (([\"] basic_string_1 [\"])) + basic_string_1 ::= "" | [^"\\\x00-\x1F] basic_string_1 | "\\" escape basic_string_1 + escape ::= ["\\/bfnrt] | "u" [A-Fa-f0-9]{4} + basic_integer ::= ("0" | "-"? [1-9] [0-9]*) ".0"? + basic_number ::= ("0" | "-"? [1-9] [0-9]*) ("." [0-9]+)? ([eE] [+-]? [0-9]+)? + basic_array ::= "[" ("" | ws basic_any (ws "," ws basic_any)*) ws "]" + basic_object ::= "{" ("" | ws basic_string ws ":" ws basic_any ( ws "," ws basic_string ws ":" ws basic_any)*) ws "}" + ws ::= [ \n\t]* + """ + + # Format-specific extensions + json_grammar_ebnf_str = ( + r""" + json ::= basic_array | basic_object + basic_any ::= basic_number | basic_string | basic_boolean | basic_null | basic_array | basic_object + basic_boolean ::= "true" | "false" + basic_null ::= "null" + """ + + BASE_PRIMITIVE_GRAMMAR + ) + + pythonic_grammar_ebnf_str = ( + r""" + pythonic ::= basic_number | basic_string | basic_array | "True" | "False" | "None" + basic_any ::= basic_number | basic_string | basic_array | basic_object + basic_boolean ::= "True" | "False" + basic_null ::= "None" + """ + + BASE_PRIMITIVE_GRAMMAR + ) + + xml_grammar_ebnf_str = ( + r""" + xml ::= xml_element | xml_text + xml_element ::= basic_string | basic_number | basic_boolean | basic_null | basic_array | basic_object + xml_text ::= [^<>]* + basic_any ::= basic_number | basic_string | basic_boolean | basic_null | basic_array | basic_object + basic_boolean ::= "true" | "false" + basic_null ::= "null" + """ + + BASE_PRIMITIVE_GRAMMAR + ) + + CALL_RULE_MAP = { + "pythonic": 'call_{name} ::= "{name}" "(" {arguments_rule} ")"', + "json": 'call_{name} ::= "{{" ws "\\"name\\"" ws ":" ws "\\"{name}\\"" ws "," ws "\\"arguments\\"" ws ":" ws {arguments_rule} ws "}}"', + "xml": 'call_{name} ::= "\\n" {arguments_rule} "\\n"', + } + + ARGUMENTS_RULE_MAP = { + "pythonic": "{arg_rules}", + "json": '"{{" ws {arg_rules} ws "}}"', + "xml": "{arg_rules}", + } + + KEY_VALUE_RULE_MAP = { + "pythonic": '"{key}" "=" {valrule}', + "json": '"\\"{key}\\"" ws ":" ws {valrule}', + "xml": '"\\n" {valrule} "\\n"', + } + + # Base type mapping - most types are the same across formats + BASE_TYPE_MAPPING = { + "string": "basic_string", + "number": "basic_number", + "integer": "basic_number", + "boolean": "basic_boolean", + "null": "basic_null", + "array": "basic_array", + "object": "basic_object", + } + + # Format-specific overrides for types that differ + FORMAT_TYPE_OVERRIDES = { + "pythonic": { + "boolean": '"True" | "False"', + "null": '"None"', + }, + "xml": { + "string": "xml_text", + }, + } + + @staticmethod + def get_value_rule( + prop: dict, function_format: Literal["pythonic", "json", "xml"] = "json" + ) -> str: + if "enum" in prop: + return EBNFComposer._handle_enum(prop, function_format) + + if "type" in prop: + return EBNFComposer._handle_type(prop, function_format) + + return function_format + + @staticmethod + def _handle_enum(prop: dict, function_format: str) -> str: + """Handle enum properties by formatting each value according to type and format.""" + enum_values = prop["enum"] + prop_type = prop.get("type", "string") + + def format_enum_val(v: Any) -> str: + if prop_type == "boolean": + if function_format == "json" or function_format == "xml": + return "true" if v else "false" + elif function_format == "pythonic": + return "True" if v else "False" + else: + return str(v) # fallback + + if prop_type == "string": + if function_format == "xml": + return f'"{v}"' + else: # json or pythonic + return f'"\\"{v}\\""' # escape quote-wrapped string + + # All other types (number, integer, etc.) + return str(v) + + formatted_values = [format_enum_val(v) for v in enum_values] + enum_rule = " | ".join(formatted_values) + return f"({enum_rule})" if len(formatted_values) > 1 else enum_rule + + @staticmethod + def get_type_mapping(function_format: str) -> Dict[str, str]: + """Get the complete type mapping for a given format.""" + mapping = EBNFComposer.BASE_TYPE_MAPPING.copy() + overrides = EBNFComposer.FORMAT_TYPE_OVERRIDES.get(function_format, {}) + mapping.update({k: v for k, v in overrides.items() if v is not None}) + return mapping + + @staticmethod + def _handle_type(prop: dict, function_format: str) -> str: + """Handle type properties using the appropriate type mapping.""" + prop_type = prop["type"] + type_mapping = EBNFComposer.get_type_mapping(function_format) + + if isinstance(prop_type, list): + type_rules = [ + type_mapping.get(single_type, function_format) + for single_type in prop_type + ] + return " | ".join(type_rules) if type_rules else function_format + + return type_mapping.get(prop_type, function_format) + + @staticmethod + def build_ebnf( + tools, + function_format: Literal["pythonic", "json", "xml"] = "json", + # Parameters for wrapping the entire sequence of tool calls + sequence_start_token: Optional[str] = None, + sequence_end_token: Optional[str] = None, + # Parameters for wrapping individual tool calls + individual_call_start_token: Optional[str] = None, + individual_call_end_token: Optional[str] = None, + # Parameter for separating multiple tool calls + tool_call_separator: Optional[str] = None, + call_rule_fmt: Optional[str] = None, + key_value_rule_fmt: Optional[str] = None, + key_value_separator: str = 'ws "," ws', + ): + """ + Generalized EBNF builder for all detectors. + Args: + tools: List of Tool objects to generate EBNF grammar for + function_format: The format of function calls, either "pythonic" or "json" + sequence_start_token: Token that wraps the entire sequence of tool calls (start) + sequence_end_token: Token that wraps the entire sequence of tool calls (end) + individual_call_start_token: Token that wraps each individual tool call (start) + individual_call_end_token: Token that wraps each individual tool call (end) + tool_call_separator: The separator between multiple tool calls + call_rule_fmt: Optional custom format string for call_{name} rule. It should define each function call's format, with + the placeholders {name} for the function name and {arguments_rule} for the arguments rule. If None, a default + format based on function_format will be used. + key_value_rule_fmt: Optional custom format string for key-value pairs. It should define how each parameter is formatted, + with placeholders {key} for the parameter name and {valrule} for the value rule. If None, a default format + based on function_format will be used. + key_value_separator: Raw EBNF fragment inserted between key-value pairs. + This string is used verbatim (not auto-quoted). Pass: + - Quoted terminals when you need a literal token (e.g. '","' or '"\\n"'). + - Raw/non-terminals when you need grammar tokens (e.g. 'ws "," ws'). + """ + # ================================================================= + # Step 1: Determine the root tool calls rule + # ================================================================= + # Handle a single function call + if individual_call_start_token and individual_call_end_token: + function_call_unit = f'"{individual_call_start_token}" function_call "{individual_call_end_token}"' + else: + function_call_unit = "function_call" + + # Handle multiple function calls with separators + if tool_call_separator is not None: + base_pattern = f'{function_call_unit} ( "{tool_call_separator}" {function_call_unit} )*' + else: + # Assume only support single function call + base_pattern = function_call_unit + + # Apply sequence-level wrapping if needed + if sequence_start_token and sequence_end_token: + root_rule = ( + f'"{sequence_start_token}" {base_pattern} "{sequence_end_token}"' + ) + else: + root_rule = base_pattern + + # ================================================================= + # Step 2: Build the header rules + # ================================================================= + ebnf_lines = [ + f"root ::= {root_rule}", + "function_call ::= " + + " | ".join([f"call_{tool.function.name}" for tool in tools]), + ] + + # ================================================================= + # Step 3: Set up formatting templates + # ================================================================= + call_template = ( + f"call_{{name}} ::= {call_rule_fmt}" + if call_rule_fmt + else EBNFComposer.CALL_RULE_MAP[function_format] + ) + args_template = EBNFComposer.ARGUMENTS_RULE_MAP[function_format] + key_value_template = ( + key_value_rule_fmt + if key_value_rule_fmt + else EBNFComposer.KEY_VALUE_RULE_MAP[function_format] + ) + + # ================================================================= + # Step 4: Build rules for each tool + # ================================================================= + for tool in tools: + tool_name = tool.function.name + params = tool.function.parameters or {} + properties = params.get("properties", {}) + required_props = set(params.get("required", [])) + + # The generated pattern ensures: + # 1. Required properties appear first, joined by commas + # 2. Optional properties are wrapped with comma included: ( "," ( "prop" : value )? )? + # 3. For multiple optional properties, we allow flexible ordering: + # - Each optional can be skipped entirely + # - They can appear in any combination + # + # Example patterns generated: + # - One required, one optional: + # "{" "location" ":" string ( "," ( "unit" ":" enum ) )? "}" + # Allows: {"location": "Paris"} or {"location": "Paris", "unit": "celsius"} + # + # - Multiple optional properties with flexible ordering: + # "{" "req" ":" string ( "," ( "opt1" ":" value ( "," "opt2" ":" value )? | "opt2" ":" value ) )? "}" + # Allows: {"req": "x"}, {"req": "x", "opt1": "y"}, {"req": "x", "opt2": "z"}, + # {"req": "x", "opt1": "y", "opt2": "z"} + # + # - All optional properties with flexible ordering: + # "{" ( "opt1" ":" value ( "," "opt2" ":" value )? | "opt2" ":" value )? "}" + # Allows: {}, {"opt1": "x"}, {"opt2": "y"}, {"opt1": "x", "opt2": "y"} + + prop_kv_pairs = {} + ordered_props = list(properties.keys()) + + for prop_name, prop_schema in properties.items(): + value_rule = EBNFComposer.get_value_rule(prop_schema, function_format) + # Create key=value pair + pair = key_value_template.format(key=prop_name, valrule=value_rule) + prop_kv_pairs[prop_name] = pair + + # Separate into required and optional while preserving order + required = [p for p in ordered_props if p in required_props] + optional = [p for p in ordered_props if p not in required_props] + + # Build the combined rule + rule_parts = [] + + # Add required properties joined by commas + if required: + rule_parts.append( + f" {key_value_separator} ".join(prop_kv_pairs[k] for k in required) + ) + + # Add optional properties with flexible ordering + if optional: + # Build alternatives where any optional property can appear first + opt_alternatives = [] + for i in range(len(optional)): + # Build pattern for optional[i] appearing first + opt_parts = [] + for j in range(i, len(optional)): + if j == i: + opt_parts.append(prop_kv_pairs[optional[j]]) + else: + opt_parts.append( + f" ( {key_value_separator} {prop_kv_pairs[optional[j]]} )?" + ) + opt_alternatives.append("".join(opt_parts)) + + # Wrap with appropriate comma handling based on whether we have required properties + if required: + # Required properties exist, so optional group needs outer comma + rule_parts.append(f" ( {key_value_separator} ( ") + rule_parts.append(" | ".join(opt_alternatives)) + rule_parts.append(" ) )?") + else: + # All properties are optional + rule_parts.append("( ") + rule_parts.append(" | ".join(opt_alternatives)) + rule_parts.append(" )?") + + combined_args = "".join(rule_parts) + arguments_rule = args_template.format(arg_rules=combined_args) + arguments_rule = arguments_rule or '""' + + # Add the function call rule and its arguments rule + ebnf_lines.append( + call_template.format( + name=tool_name, arguments_rule=f"arguments_{tool_name}" + ) + ) + ebnf_lines.append(f"arguments_{tool_name} ::= {arguments_rule}") + + # ================================================================= + # Step 5: Add base grammar rules + # ================================================================= + grammar_dict = { + "pythonic": EBNFComposer.pythonic_grammar_ebnf_str, + "json": EBNFComposer.json_grammar_ebnf_str, + "xml": EBNFComposer.xml_grammar_ebnf_str, + } + base_grammar = grammar_dict.get( + function_format, EBNFComposer.json_grammar_ebnf_str + ) + ebnf_lines.append(base_grammar) + + return "\n".join(ebnf_lines) diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py new file mode 100644 index 000000000..18fe488e4 --- /dev/null +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -0,0 +1,214 @@ +import logging +from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Type, Union + +from sglang.srt.entrypoints.openai.protocol import ( + StructuralTagResponseFormat, + StructuresResponseFormat, + Tool, + ToolChoice, +) +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ToolCallItem +from sglang.srt.function_call.deepseekv3_detector import DeepSeekV3Detector +from sglang.srt.function_call.deepseekv31_detector import DeepSeekV31Detector +from sglang.srt.function_call.glm4_moe_detector import Glm4MoeDetector +from sglang.srt.function_call.gpt_oss_detector import GptOssDetector +from sglang.srt.function_call.kimik2_detector import KimiK2Detector +from sglang.srt.function_call.llama32_detector import Llama32Detector +from sglang.srt.function_call.mistral_detector import MistralDetector +from sglang.srt.function_call.pythonic_detector import PythonicDetector +from sglang.srt.function_call.qwen3_coder_detector import Qwen3CoderDetector +from sglang.srt.function_call.qwen25_detector import Qwen25Detector +from sglang.srt.function_call.step3_detector import Step3Detector + +logger = logging.getLogger(__name__) + + +class FunctionCallParser: + """ + Parser for function/tool calls in model outputs. + + This class handles both streaming and non-streaming parsing of function calls using a detector. + In streaming scenarios, each time new_text is received, it calls detector.parse_streaming_increment + and returns the resulting normal_text and calls to the upper layer (or SSE). + """ + + ToolCallParserEnum: Dict[str, Type[BaseFormatDetector]] = { + "llama3": Llama32Detector, + "qwen25": Qwen25Detector, + "mistral": MistralDetector, + "deepseekv3": DeepSeekV3Detector, + "deepseekv31": DeepSeekV31Detector, + "pythonic": PythonicDetector, + "kimi_k2": KimiK2Detector, + "qwen3_coder": Qwen3CoderDetector, + "glm45": Glm4MoeDetector, + "step3": Step3Detector, + "gpt-oss": GptOssDetector, + } + + def __init__(self, tools: List[Tool], tool_call_parser: str): + detector: Type[BaseFormatDetector] = None + detector_class = self.ToolCallParserEnum.get(tool_call_parser) + if detector_class: + detector = detector_class() + else: + raise ValueError(f"Unsupported tool_call_parser: {tool_call_parser}") + + self.detector = detector + self.tools = tools + + def has_tool_call(self, text: str) -> bool: + """ + Check if the given text contains a tool call in the format supported by this parser. + This delegates to the detector's implementation. + + Args: + text: The text to check for tool calls + + Returns: + True if the text contains a tool call, False otherwise + """ + return self.detector.has_tool_call(text) + + def parse_non_stream(self, full_text: str) -> Tuple[str, list[ToolCallItem]]: + """ + One-time parsing of the full text to extract tool calls. + + Args: + full_text: The complete text to parse + + Returns: + A tuple containing: + - The remaining text after parsing that was not consumed by the detector (can be treated as normal text) + - A list of tool calls parsed from the text + """ + parsed_result = self.detector.detect_and_parse(full_text, self.tools) + tool_call_list = parsed_result.calls + if tool_call_list: + return parsed_result.normal_text, tool_call_list + else: + return full_text, [] + + def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, list[ToolCallItem]]: + """ + Streaming incremental parsing of chunks of text as they arrive. + + Args: + chunk_text: The new chunk of text to parse + + Returns: + A tuple containing: + - The normal text that should be displayed to the user + - A list of tool calls parsed from the chunk + """ + final_normal_text = "" + final_calls = [] + + sp_result = self.detector.parse_streaming_increment(chunk_text, self.tools) + if sp_result.normal_text: + final_normal_text = sp_result.normal_text + if sp_result.calls: + final_calls.extend(sp_result.calls) + final_normal_text = sp_result.normal_text + + return final_normal_text, final_calls + + def get_structure_tag(self) -> StructuralTagResponseFormat: + """ + Generate a structural tag response format for all available tools. + + This creates the necessary structural tags that guide the model's output format. + """ + tool_structures: List[StructuresResponseFormat] = list() + tool_trigger_set: Set[str] = set() + + get_structure_info = self.detector.structure_info() + for tool in self.tools: + function = tool.function + name = function.name + assert name is not None + info = get_structure_info(name) + + # accept all if not strict, otherwise only accept the schema + schema = function.parameters if function.strict else {} + + tool_structures.append( + StructuresResponseFormat( + begin=info.begin, + schema=schema, # type: ignore + end=info.end, + ) + ) + tool_trigger_set.add(info.trigger) + + return StructuralTagResponseFormat( + type="structural_tag", + structures=tool_structures, + triggers=list(tool_trigger_set), + ) + + def get_structure_constraint( + self, tool_choice: Union[ToolChoice, Literal["auto", "required"]] + ) -> Optional[Tuple[str, Any]]: + """ + Returns the appropriate structure constraint for tool calls based on the tool_choice. + The constraint is used to guide the model's output format. + + Args: + tool_choice: The tool choice setting from the request + + Returns: + A tuple of (constraint_type, constraint_value) to be added to sampling parameters, + or None if no constraint applies. + """ + # NOTE: structural_tag only supports JSON-compatible content between the begin and end. + # It cannot parse or validate function call Pythonic or XML-ish syntax. + if ( + self.detector.supports_structural_tag() + and tool_choice == "auto" + and any(tool.function.strict for tool in self.tools) + ): + strict_tag = self.get_structure_tag() + return ("structural_tag", strict_tag) + elif tool_choice == "required" or isinstance(tool_choice, ToolChoice): + ebnf = self.get_ebnf(tool_choice) + return ("ebnf", ebnf) if ebnf is not None else None + + def get_ebnf( + self, tool_choice: Union[ToolChoice, Literal["required"]] + ) -> Optional[str]: + """ + Get the EBNF grammar for the specified tool choice. + + Args: + tool_choice: The tool choice specification + + Returns: + EBNF grammar string, or None if no valid tools found + + Note: + If a specific function is requested but not found in available tools, + logs a warning and falls back to using all available tools for backward compatibility. + """ + filtered_tools = [] + if isinstance(tool_choice, ToolChoice): + fn_name = tool_choice.function.name + filtered_tools = [t for t in self.tools if t.function.name == fn_name] + + # Check if the requested function exists in available tools + if not filtered_tools: + available_functions = [t.function.name for t in self.tools] + logger.warning( + f"Function '{fn_name}' not found in available tools. " + f"Available functions: {available_functions}. " + f"Skipping tool choice." + ) + + # TODO: Return a 400 error instead of warning when adapter supports proper error handling + # For now, fall back to return None + return None + else: + filtered_tools = self.tools + + return self.detector.build_ebnf(filtered_tools) diff --git a/python/sglang/srt/function_call/glm4_moe_detector.py b/python/sglang/srt/function_call/glm4_moe_detector.py new file mode 100644 index 000000000..6e89fe0a1 --- /dev/null +++ b/python/sglang/srt/function_call/glm4_moe_detector.py @@ -0,0 +1,164 @@ +import ast +import json +import logging +import re +from typing import List + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer + +logger = logging.getLogger(__name__) + + +def get_argument_type(func_name: str, arg_key: str, defined_tools: list): + name2tool = {tool.function.name: tool for tool in defined_tools} + if func_name not in name2tool: + return None + tool = name2tool[func_name] + if arg_key not in tool.function.parameters["properties"]: + return None + return tool.function.parameters["properties"][arg_key].get("type", None) + + +def parse_arguments(json_value): + try: + try: + parsed_value = json.loads(json_value) + except: + parsed_value = ast.literal_eval(json_value) + return parsed_value, True + except: + return json_value, False + + +class Glm4MoeDetector(BaseFormatDetector): + """ + Detector for GLM-4.5 models. + Assumes function call format: + get_weather\ncity\n北京\ndate\n2024-06-27\n\nget_weather\ncity\n上海\ndate\n2024-06-27\n + """ + + def __init__(self): + super().__init__() + self.bot_token = "" + self.eot_token = "" + self.func_call_regex = r".*?" + self.func_detail_regex = r"([^\n]*)\n(.*)" + self.func_arg_regex = r"(.*?)\s*(.*?)" + + def has_tool_call(self, text: str) -> bool: + """Check if the text contains a glm-4.5 format tool call.""" + return self.bot_token in text + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """ + One-time parsing: Detects and parses tool calls in the provided text. + + :param text: The complete text to parse. + :param tools: List of available tools. + :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. + """ + idx = text.find(self.bot_token) + normal_text = text[:idx].strip() if idx != -1 else text + if self.bot_token not in text: + return StreamingParseResult(normal_text=normal_text, calls=[]) + match_result_list = re.findall(self.func_call_regex, text, re.DOTALL) + calls = [] + try: + for match_result in match_result_list: + # Get function name + func_detail = re.search(self.func_detail_regex, match_result, re.DOTALL) + func_name = func_detail.group(1) + func_args = func_detail.group(2) + pairs = re.findall( + r"(.*?)\s*(.*?)", + func_args, + re.DOTALL, + ) + arguments = {} + for arg_key, arg_value in pairs: + arg_key = arg_key.strip() + arg_value = arg_value.strip() + arg_type = get_argument_type(func_name, arg_key, tools) + if arg_type != "string": + arg_value, is_good_json = parse_arguments(arg_value) + arguments[arg_key] = arg_value + # construct match_result for parse_base_json + match_result = {"name": func_name, "parameters": arguments} + calls.extend(self.parse_base_json(match_result, tools)) + return StreamingParseResult(normal_text=normal_text, calls=calls) + except Exception as e: + logger.error(f"Error in detect_and_parse: {e}") + # return the normal text if parsing fails + return StreamingParseResult(normal_text=text) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Streaming incremental parsing tool calls for GLM-4.5 format. + """ + self._buffer += new_text + current_text = self._buffer + + start = current_text.find(self.bot_token) + if start == -1: + self._buffer = "" + if self.current_tool_id > 0: + current_text = "" + return StreamingParseResult(normal_text=current_text) + # find ensures we find the first self.eot_token so there will be at most one tool_call in current_text[:end+len(self.eot_token) + end = current_text.find(self.eot_token) + if end != -1: + # Initialize state if this is the first tool call + if self.current_tool_id == -1: + self.current_tool_id = 0 + self.prev_tool_call_arr = [] + self.streamed_args_for_tool = [""] + # Ensure we have enough entries in our tracking arrays + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + result = self.detect_and_parse( + current_text[: end + len(self.eot_token)], tools=tools + ) + if result.calls: + self.prev_tool_call_arr[self.current_tool_id] = { + "name": result.calls[0].name, + "arguments": json.loads(result.calls[0].parameters), + } + self.streamed_args_for_tool[self.current_tool_id] = result.calls[ + 0 + ].parameters + result.calls[0].tool_index = self.current_tool_id + self.current_tool_id += 1 + self._buffer = current_text[end + len(self.eot_token) :] + return result + normal_text = current_text[:start] + self._buffer = current_text[start:] + return StreamingParseResult(normal_text=normal_text) + + def supports_structural_tag(self) -> bool: + return False + + def structure_info(self) -> _GetInfoFunc: + raise NotImplementedError() + + def build_ebnf(self, tools: List[Tool]): + return EBNFComposer.build_ebnf( + tools, + individual_call_start_token=self.bot_token, + individual_call_end_token=self.eot_token, + tool_call_separator="\\n", + function_format="xml", + call_rule_fmt='"{name}" "\\n" ( {arguments_rule} "\\n" )?', + key_value_rule_fmt='"{key}" "\\n" "" {valrule} ""', + key_value_separator='"\\n"', + ) diff --git a/python/sglang/srt/function_call/gpt_oss_detector.py b/python/sglang/srt/function_call/gpt_oss_detector.py new file mode 100644 index 000000000..7fe0a7dc8 --- /dev/null +++ b/python/sglang/srt/function_call/gpt_oss_detector.py @@ -0,0 +1,219 @@ +import json +import logging +import re +from typing import List, Optional + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.parser.harmony_parser import HarmonyParser + +logger = logging.getLogger(__name__) + + +class GptOssDetector(BaseFormatDetector): + """ + Detector for T4-style function calls using HarmonyParser. + + Handles tool calls in the format: + <|channel|>commentary to={namespace.function}<|constrain|>json<|message|>{args}<|call|> + """ + + def __init__(self): + super().__init__() + self.harmony_parser = HarmonyParser() + self.bot_token = "<|start|>assistant<|channel|>commentary" + self.eot_token = "<|call|>" + + # Pattern to extract function name and JSON from tool_call event content + self.tool_extract_pattern = re.compile( + r"to=([a-zA-Z_][a-zA-Z0-9_.]*)\s*<\|constrain\|>json<\|message\|>(.*?)(?:<\|call\|>|$)", + re.DOTALL, + ) + + def has_tool_call(self, text: str) -> bool: + """Check if text contains TypeScript-style function call markers.""" + return self.bot_token in text + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """Parse TypeScript-style function calls from complete text.""" + if not self.has_tool_call(text): + return StreamingParseResult(normal_text=text, calls=[]) + + # Parse with HarmonyParser + events = self.harmony_parser.parse(text) + # Flush buffer for complete parsing + events += self.harmony_parser.parse("") + + tool_indices = self._get_tool_indices(tools) + calls = [] + normal_parts = [] + tool_index = 0 + + for event in events: + if event.event_type == "tool_call": + # Extract tool call from event content + tool_call = self._extract_tool_call_from_event( + event.raw_text if event.raw_text else event.content, + tool_indices, + tool_index, + ) + if tool_call: + calls.append(tool_call) + tool_index += 1 + elif event.event_type == "normal": + normal_parts.append(event.content) + # Ignore reasoning events in function call context + + normal_text = " ".join(normal_parts).strip() + return StreamingParseResult(normal_text=normal_text, calls=calls) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """Parse incremental streaming text for TypeScript-style function calls.""" + self._buffer += new_text + + # Always use HarmonyParser for parsing to ensure proper filtering + events = self.harmony_parser.parse(new_text) + + # Quick check if we might have tool calls + if ( + "<|channel|>commentary to=" not in self._buffer + and not self.current_tool_name_sent + ): + # No tool calls detected, check for final content + if ( + "<|channel|>final" in self._buffer + or "assistantfinal" in self._buffer.lower() + ): + # Extract normal text from events + normal_text = "".join( + [e.content for e in events if e.event_type == "normal"] + ) + if normal_text: + self._buffer = "" + return StreamingParseResult(normal_text=normal_text, calls=[]) + + # For other content, extract normal text from events (with filtering applied) + normal_text = "".join( + [e.content for e in events if e.event_type == "normal"] + ) + if normal_text or events: + self._buffer = "" + return StreamingParseResult(normal_text=normal_text, calls=[]) + else: + # No events processed, continue buffering + return StreamingParseResult(normal_text="", calls=[]) + + if not events: + # No complete events yet + return StreamingParseResult(normal_text="", calls=[]) + + # Initialize state if needed + if not hasattr(self, "_tool_indices"): + self._tool_indices = self._get_tool_indices(tools) + + calls = [] + normal_text = "" + + for event in events: + if event.event_type == "tool_call": + # We got a complete tool call from HarmonyParser + tool_call_info = self._extract_tool_call_from_event( + event.raw_text if event.raw_text else event.content, + self._tool_indices, + self.current_tool_id if self.current_tool_id >= 0 else 0, + ) + + if tool_call_info: + # Initialize state if first tool + if self.current_tool_id == -1: + self.current_tool_id = 0 + self.prev_tool_call_arr = [] + self.streamed_args_for_tool = [""] + + # Ensure arrays are large enough + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + + # Store tool call info + self.prev_tool_call_arr[self.current_tool_id] = { + "name": tool_call_info.name, + "arguments": json.loads(tool_call_info.parameters), + } + + # Emit the complete tool call at once + # (Could be modified to emit name first, then args, if needed) + calls.append(tool_call_info) + + # Mark as streamed + self.streamed_args_for_tool[self.current_tool_id] = ( + tool_call_info.parameters + ) + + # Move to next tool + self.current_tool_id += 1 + self.current_tool_name_sent = False + + elif event.event_type == "normal": + normal_text += event.content + + # Clear buffer since HarmonyParser handles buffering + self._buffer = "" + + return StreamingParseResult(normal_text=normal_text, calls=calls) + + def _extract_tool_call_from_event( + self, content: str, tool_indices: dict, tool_index: int + ) -> Optional[ToolCallItem]: + """ + Extract tool call information from HarmonyParser event content. + + Content format: "commentary to=functions.get_weather<|constrain|>json<|message|>{...}" + """ + match = self.tool_extract_pattern.search(content) + + if not match: + logger.debug(f"Could not extract tool call from: {content[:100]}") + return None + + full_function_name = match.group(1) + json_content = match.group(2) + + # Extract function name (last part after .) + function_name = ( + full_function_name.split(".")[-1] + if "." in full_function_name + else full_function_name + ) + + # Check if tool exists + if function_name not in tool_indices: + logger.debug(f"Function {function_name} not in available tools") + return None + + # Parse JSON arguments + try: + arguments = json.loads(json_content) if json_content.strip() else {} + except json.JSONDecodeError as e: + logger.debug(f"Failed to parse JSON arguments: {e}") + return None + + return ToolCallItem( + tool_index=tool_index, + name=function_name, + parameters=json.dumps(arguments, ensure_ascii=False), + ) + + def structure_info(self) -> _GetInfoFunc: + raise NotImplementedError("structure_info not used with HarmonyParser") + + def build_ebnf(self, tools: List[Tool]) -> str: + raise NotImplementedError("build_ebnf not used with HarmonyParser") diff --git a/python/sglang/srt/function_call/kimik2_detector.py b/python/sglang/srt/function_call/kimik2_detector.py new file mode 100644 index 000000000..4f39433a6 --- /dev/null +++ b/python/sglang/srt/function_call/kimik2_detector.py @@ -0,0 +1,245 @@ +import json +import logging +import re +from typing import List + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer +from sglang.srt.function_call.utils import _is_complete_json + +logger = logging.getLogger(__name__) + + +class KimiK2Detector(BaseFormatDetector): + """ + Detector for Kimi K2 model function call format. + + Format Structure: + ``` + <|tool_calls_section_begin|> + <|tool_call_begin|>functions.{func_name}:{index}<|tool_call_argument_begin|>{json_args}<|tool_call_end|> + <|tool_calls_section_end|> + ``` + + Reference: https://huggingface.co/moonshotai/Kimi-K2-Instruct/blob/main/docs/tool_call_guidance.md + """ + + def __init__(self): + super().__init__() + + self.bot_token: str = "<|tool_calls_section_begin|>" + self.eot_token: str = "<|tool_calls_section_end|>" + + self.tool_call_start_token: str = "<|tool_call_begin|>" + self.tool_call_end_token: str = "<|tool_call_end|>" + + self.tool_call_regex = re.compile( + r"<\|tool_call_begin\|>\s*(?P[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P\{.*?\})\s*<\|tool_call_end\|>" + ) + + self.stream_tool_call_portion_regex = re.compile( + r"<\|tool_call_begin\|>\s*(?P[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P\{.*)" + ) + + self._last_arguments = "" + + def has_tool_call(self, text: str) -> bool: + """Check if the text contains a KimiK2 format tool call.""" + return self.bot_token in text + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """ + One-time parsing: Detects and parses tool calls in the provided text. + + :param text: The complete text to parse. + :param tools: List of available tools. + :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. + """ + if self.bot_token not in text: + return StreamingParseResult(normal_text=text, calls=[]) + try: + # there are two possible captures - between tags, or between a + # tag and end-of-string so the result of + # findall is an array of tuples where one is a function call and + # the other is None + function_call_tuples = self.tool_call_regex.findall(text) + + logger.debug("function_call_tuples: %s", function_call_tuples) + + tool_calls = [] + for match in function_call_tuples: + function_id, function_args = match + function_name = function_id.split(".")[1].split(":")[0] + function_idx = int(function_id.split(".")[1].split(":")[1]) + + logger.info(f"function_name {function_name}") + + tool_calls.append( + ToolCallItem( + tool_index=function_idx, # Use the call index in the response, not tool position + name=function_name, + parameters=function_args, + ) + ) + + content = text[: text.find(self.bot_token)] + return StreamingParseResult(normal_text=content, calls=tool_calls) + + except Exception as e: + logger.error(f"Error in detect_and_parse: {e}") + # return the normal text if parsing fails + return StreamingParseResult(normal_text=text) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Streaming incremental parsing tool calls for KimiK2 format. + """ + self._buffer += new_text + current_text = self._buffer + + # Check if we have a tool call (either the start token or individual tool call) + has_tool_call = ( + self.bot_token in current_text or self.tool_call_start_token in current_text + ) + + if not has_tool_call: + self._buffer = "" + for e_token in [self.eot_token, self.tool_call_end_token]: + if e_token in new_text: + new_text = new_text.replace(e_token, "") + return StreamingParseResult(normal_text=new_text) + + if not hasattr(self, "_tool_indices"): + self._tool_indices = self._get_tool_indices(tools) + + calls: list[ToolCallItem] = [] + try: + match = self.stream_tool_call_portion_regex.search(current_text) + if match: + function_id = match.group("tool_call_id") + function_args = match.group("function_arguments") + + function_name = function_id.split(".")[1].split(":")[0] + + # Initialize state if this is the first tool call + if self.current_tool_id == -1: + self.current_tool_id = 0 + self.prev_tool_call_arr = [] + self.streamed_args_for_tool = [""] + + # Ensure we have enough entries in our tracking arrays + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + + if not self.current_tool_name_sent: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=function_name, + parameters="", + ) + ) + self.current_tool_name_sent = True + # Store the tool call info for serving layer completions endpoint + self.prev_tool_call_arr[self.current_tool_id] = { + "name": function_name, + "arguments": {}, + } + else: + argument_diff = ( + function_args[len(self._last_arguments) :] + if function_args.startswith(self._last_arguments) + else function_args + ) + + parsed_args_diff = argument_diff.split("<|tool_call_end|>", 1)[0] + + if parsed_args_diff: + + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=None, + parameters=parsed_args_diff, + ) + ) + self._last_arguments += argument_diff + self.streamed_args_for_tool[ + self.current_tool_id + ] += parsed_args_diff + + parsed_args = function_args.split("<|tool_call_end|>", 1)[0] + if _is_complete_json(parsed_args): + try: + parsed_args = json.loads(parsed_args) + self.prev_tool_call_arr[self.current_tool_id][ + "arguments" + ] = parsed_args + except json.JSONDecodeError: + pass + + # Find the end of the current tool call and remove only that part from buffer + tool_call_end_pattern = ( + r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>" + ) + match = re.search( + tool_call_end_pattern, current_text, re.DOTALL + ) + if match: + # Remove the completed tool call from buffer, keep any remaining content + self._buffer = current_text[match.end() :] + else: + self._buffer = "" + + result = StreamingParseResult(normal_text="", calls=calls) + self.current_tool_id += 1 + self._last_arguments = "" + self.current_tool_name_sent = False + return result + + return StreamingParseResult(normal_text="", calls=calls) + + except Exception as e: + logger.error(f"Error in parse_streaming_increment: {e}") + return StreamingParseResult(normal_text=current_text) + + def structure_info(self) -> _GetInfoFunc: + """Return function that creates StructureInfo for guided generation.""" + + def get_info(name: str) -> StructureInfo: + return StructureInfo( + begin=f"<|tool_calls_section_begin|><|tool_call_begin|>functions.{name}:0<|tool_call_argument_begin|>", + end="<|tool_call_end|><|tool_calls_section_end|>", + trigger="<|tool_calls_section_begin|>", + ) + + return get_info + + def build_ebnf(self, tools: List[Tool]) -> str: + """ + Build EBNF grammar for KimiK2 tool call format. + + NOTE: The call_rule_fmt uses [0-9]+ for the function index to allow the grammar + to accept any numeric index (0, 1, 2, etc.) for proper sequential indexing in + multiple function call scenarios, while still maintaining the correct KimiK2 + format structure for constrained generation. + """ + return EBNFComposer.build_ebnf( + tools, + sequence_start_token=self.bot_token, + sequence_end_token=self.eot_token, + tool_call_separator="", + call_rule_fmt='"<|tool_call_begin|>functions.{name}:"[0-9]+"<|tool_call_argument_begin|>"{arguments_rule}"<|tool_call_end|>"', + function_format="json", + ) diff --git a/python/sglang/srt/function_call/llama32_detector.py b/python/sglang/srt/function_call/llama32_detector.py new file mode 100644 index 000000000..453bcbc9a --- /dev/null +++ b/python/sglang/srt/function_call/llama32_detector.py @@ -0,0 +1,96 @@ +import json +import logging +from typing import List + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer + +logger = logging.getLogger(__name__) + + +class Llama32Detector(BaseFormatDetector): + """ + Detector for Llama 3.2 models with json tool call format. + + Format Structure: + ``` + {"name":"xxx", "arguments":{...}} + ``` + """ + + def __init__(self): + super().__init__() + self.bot_token = "<|python_tag|>" + # NOTE: technically Llama3.2 doesn't support well with parallel tool calls + # They need specific prompt engineering to support parallel tool calls + # Here we use ';' as the separator, which might have compatibility issues + # if users define to use a different separator in their prompt + self.tool_call_separator = ";" + + def has_tool_call(self, text: str) -> bool: + """Check if the text contains a Llama 3.2 format tool call.""" + # depending on the prompt format the Llama model may or may not + # prefix the output with the <|python_tag|> token + return "<|python_tag|>" in text or text.startswith("{") + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """Parse function calls from text, handling multiple JSON objects.""" + if "<|python_tag|>" not in text and not text.startswith("{"): + return StreamingParseResult(normal_text=text, calls=[]) + + if "<|python_tag|>" in text: + normal_text, action_text = text.split("<|python_tag|>", maxsplit=1) + else: + normal_text, action_text = "", text + + decoder = json.JSONDecoder() + idx = 0 + safe_idx = idx # the index of the last valid JSON object + all_actions = [] + action_text_len = len(action_text) + while idx < action_text_len: + try: + obj, end = decoder.raw_decode(action_text[idx:]) + all_actions.append(obj) + idx += end + len(self.tool_call_separator) + safe_idx = idx + except json.JSONDecodeError as e: + # Find where next `{"name"` appears and try again + logger.warning( + f"Failed to parse JSON part: {action_text[idx:]}, JSON parse error: {str(e)}" + ) + next_obj_start = action_text.find('{"name":', idx + 1) + if next_obj_start == -1: + break + idx = next_obj_start + continue + + # Only process if we found valid JSON objects + calls = self.parse_base_json(all_actions, tools) if all_actions else [] + # Use safe_idx to avoid idx containing the last part of an invalid JSON object + trailing_text = ( + action_text[safe_idx:].strip() if safe_idx < action_text_len else "" + ) + return StreamingParseResult( + normal_text=normal_text + trailing_text, calls=calls + ) + + def structure_info(self) -> _GetInfoFunc: + return lambda name: StructureInfo( + begin='<|python_tag|>{"name":"' + name + '", "arguments":', + end="}", + trigger="<|python_tag|>", + ) + + def build_ebnf(self, tools: List[Tool]): + return EBNFComposer.build_ebnf( + tools, + function_format="json", + tool_call_separator=self.tool_call_separator, + ) diff --git a/python/sglang/srt/function_call/mistral_detector.py b/python/sglang/srt/function_call/mistral_detector.py new file mode 100644 index 000000000..49767fd53 --- /dev/null +++ b/python/sglang/srt/function_call/mistral_detector.py @@ -0,0 +1,139 @@ +import json +import logging +import re +from typing import List + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer + +logger = logging.getLogger(__name__) + + +class MistralDetector(BaseFormatDetector): + """ + Detector for Mistral model function call format. + + The Mistral format uses a simple bracket-delimited structure with JSON arrays + containing function call objects. + + Format Structure: + ``` + [TOOL_CALLS] [{"name": "function_name", "arguments": {json_args}}, ...] + ``` + + Reference: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3?chat_template=default + """ + + def __init__(self): + """ + Initializes the detector with necessary state variables. + """ + super().__init__() + self.bot_token = "[TOOL_CALLS] [" + self.eot_token = "]" + self.tool_call_regex = re.compile(r"\[{.*}\]", re.DOTALL) + self.tool_call_separator = ", " + + def has_tool_call(self, text: str) -> bool: + """Check if the text contains a Mistral format tool call.""" + return self.bot_token in text + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """ + One-time parsing: Detects and parses tool calls in the provided text. + + :param text: The complete text to parse. + :param tools: List of available tools. + :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. + """ + idx = text.find(self.bot_token) + normal_text = text[:idx].strip() if idx != -1 else text + + if self.bot_token not in text: + return StreamingParseResult(normal_text=normal_text, calls=[]) + + # Extract the JSON array part from [TOOL_CALLS] [...] + # Use bracket counting to properly handle nested brackets in JSON content + json_array_str = self._extract_json_array(text) + if not json_array_str: + return StreamingParseResult(normal_text=normal_text, calls=[]) + + calls = [] + try: + function_call_arr = json.loads(json_array_str) + # Handle both single object and array of objects + if not isinstance(function_call_arr, list): + function_call_arr = [function_call_arr] + calls = self.parse_base_json(function_call_arr, tools) + except json.JSONDecodeError as e: + logger.warning( + f"Failed to parse JSON part: {json_array_str}, JSON parse error: {str(e)}" + ) + + return StreamingParseResult(normal_text=normal_text, calls=calls) + + def _extract_json_array(self, text: str) -> str: + """ + Extract the JSON array part using bracket counting to handle nested brackets. + + :param text: The complete text containing [TOOL_CALLS] [...] + :return: The JSON array string or None if not found + """ + start_idx = text.find(self.bot_token) + if start_idx == -1: + return None + + # Start from the opening bracket after [TOOL_CALLS] + json_start = ( + start_idx + len(self.bot_token) - 1 + ) # -1 to include the opening bracket + bracket_count = 0 + in_string = False + escape_next = False + + for i in range(json_start, len(text)): + char = text[i] + + if escape_next: + escape_next = False + continue + + if char == "\\": + escape_next = True + continue + + if char == '"' and not escape_next: + in_string = not in_string + continue + + if not in_string: + if char == "[": + bracket_count += 1 + elif char == "]": + bracket_count -= 1 + if bracket_count == 0: + return text[json_start : i + 1] + + return None + + def structure_info(self) -> _GetInfoFunc: + return lambda name: StructureInfo( + begin='[TOOL_CALLS] [{"name":"' + name + '", "arguments":', + end="}]", + trigger="[TOOL_CALLS]", + ) + + def build_ebnf(self, tools: List[Tool]): + return EBNFComposer.build_ebnf( + tools, + sequence_start_token=self.bot_token, + sequence_end_token=self.eot_token, + function_format="json", + tool_call_separator=self.tool_call_separator, + ) diff --git a/python/sglang/srt/function_call/pythonic_detector.py b/python/sglang/srt/function_call/pythonic_detector.py new file mode 100644 index 000000000..be183c6bf --- /dev/null +++ b/python/sglang/srt/function_call/pythonic_detector.py @@ -0,0 +1,231 @@ +import ast +import json +import logging +import re +from typing import List, Optional + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer + +logger = logging.getLogger(__name__) + + +class PythonicDetector(BaseFormatDetector): + """ + Detector for Llama-4 models with Pythonic tool call format. + + The Pythonic format uses Python function call syntax within square brackets, + with arguments as Python literals rather than JSON. + + Format Structure: + ``` + [tool1(arg1=val1, arg2=val2), tool2(arg1=val3)] + ``` + + Reference: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct?chat_template=default + """ + + def __init__(self): + super().__init__() + self.tool_call_regex = re.compile( + r"\[([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*,\s*)*([a-zA-Z]+\w*=.*\s)?\),\s*)*([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*,\s*)*([a-zA-Z]+\w*=.*\s*)?\)\s*)+\]", + re.DOTALL, + ) + + @staticmethod + def _text_strip(text: str) -> str: + # Llama 4 model sometime will output <|python_start|> and <|python_end|> tokens + # remove those tokens + text = text.replace("<|python_start|>", "") + text = text.replace("<|python_end|>", "") + return text + + def has_tool_call(self, text: str) -> bool: + return bool(self.tool_call_regex.search(self._text_strip(text.strip()))) + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + # Try parsing the text as a Python list of function calls + text = text.strip() + + # Remove unexpected <|python_start|> and <|python_end|> for llama4 + text = self._text_strip(text) + + match = self.tool_call_regex.search(text) + if match is None: + return StreamingParseResult(normal_text=text, calls=[]) + + # Extract the tool call part and any text before/after it + tool_call_start = match.start() + tool_call_end = match.end() + + normal_text_before = text[:tool_call_start] if tool_call_start > 0 else "" + tool_call_text = text[tool_call_start:tool_call_end] + normal_text_after = text[tool_call_end:] if tool_call_end < len(text) else "" + + # Combine normal text + normal_text = normal_text_before + normal_text_after + + try: + module = ast.parse(tool_call_text) + parsed = getattr(module.body[0], "value", None) + if not ( + isinstance(parsed, ast.List) + and all(isinstance(e, ast.Call) for e in parsed.elts) + ): + return StreamingParseResult(normal_text=normal_text, calls=[]) + + calls = [] + tool_indices = self._get_tool_indices(tools) + for call_index, call in enumerate(parsed.elts): + if not isinstance(call.func, ast.Name): + continue + function_name = call.func.id + # Validate that the function exists in the tools + if function_name not in tool_indices: + logger.warning( + f"Model attempted to call undefined function: {function_name}" + ) + continue + arguments = {} + for keyword in call.keywords: + arguments[keyword.arg] = self._get_parameter_value(keyword.value) + calls.append( + ToolCallItem( + tool_index=call_index, # Use the call index in the response, not tool position + name=function_name, + parameters=json.dumps(arguments, ensure_ascii=False), + ) + ) + + return StreamingParseResult(normal_text=normal_text, calls=calls) + except Exception: + logger.exception("Error in pythonic tool call parsing.") + return StreamingParseResult(normal_text=normal_text, calls=[]) + + def _find_matching_bracket(self, buffer: str, start: int) -> int: + """ + Find the matching closing bracket for the opening bracket at start position. + Properly handles nested brackets. + + Args: + buffer: The text buffer to search in + start: Position of the opening bracket '[' + + Returns: + Position of the matching closing bracket ']', or -1 if not found + """ + bracket_count = 0 + for i in range(start, len(buffer)): + if buffer[i] == "[": + bracket_count += 1 + elif buffer[i] == "]": + bracket_count -= 1 + if bracket_count == 0: + return i + return -1 # No matching bracket found + + def _strip_and_split_buffer(self, buffer: str) -> tuple[str, str]: + """ + Strip special tokens from buffer and split into safe_text and held_back_text. + + Returns: + tuple of (safe_text_to_output, text_to_hold_in_buffer) + """ + # Check if original buffer ends with a partial token at the end + special_tokens = ["<|python_start|>", "<|python_end|>"] + + for token in special_tokens: + partial_length = self._ends_with_partial_token(buffer, token) + if partial_length > 0: + # Split buffer: safe part + held back partial token + safe_text = buffer[:-partial_length] + held_back = buffer[-partial_length:] + # Strip complete special tokens from safe part only + safe_text = self._text_strip(safe_text) + return safe_text, held_back + + # No partial tokens found, strip complete tokens from entire buffer + safe_text = self._text_strip(buffer) + return safe_text, "" + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Streaming incremental parsing for pythonic tool calls. + Buffers input until a complete pythonic tool call (from [ to ]) is found, + then parses and emits any detected calls. + """ + self._buffer += new_text + + # Strip special tokens from entire buffer and handle partial tokens + stripped_buffer, held_back = self._strip_and_split_buffer(self._buffer) + + start = stripped_buffer.find("[") + + if start == -1: + # No tool call bracket found + self._buffer = held_back + return StreamingParseResult(normal_text=stripped_buffer) + + normal_text = stripped_buffer[:start] if start > 0 else "" + + end = self._find_matching_bracket(stripped_buffer, start) + if end != -1: + # Found complete tool call + call_text = stripped_buffer[start : end + 1] + result = self.detect_and_parse(call_text, tools) + + # Update buffer with remaining text after tool call plus any held back text + remaining_text = stripped_buffer[end + 1 :] + held_back + self._buffer = remaining_text + + # If we had normal text before the tool call, add it to the result + if normal_text: + result.normal_text = normal_text + (result.normal_text or "") + + return result + + # We have an opening bracket but no closing bracket yet + # Put back everything from the bracket onwards plus held back text + self._buffer = stripped_buffer[start:] + held_back + + if normal_text: + return StreamingParseResult(normal_text=normal_text) + + # Otherwise, we're still accumulating a potential tool call + return StreamingParseResult(normal_text="") + + def _get_parameter_value(self, val): + if isinstance(val, ast.Constant): + return val.value + elif isinstance(val, ast.Dict): + return { + k.value: self._get_parameter_value(v) + for k, v in zip(val.keys, val.values) + } + elif isinstance(val, ast.List): + return [self._get_parameter_value(v) for v in val.elts] + else: + raise ValueError("Tool call arguments must be literals") + + def supports_structural_tag(self) -> bool: + return False + + def structure_info(self) -> _GetInfoFunc: + raise NotImplementedError + + def build_ebnf(self, tools: List[Tool]) -> Optional[str]: + return EBNFComposer.build_ebnf( + tools, + sequence_start_token="[", + sequence_end_token="]", + tool_call_separator=",", + function_format="pythonic", + ) diff --git a/python/sglang/srt/function_call/qwen25_detector.py b/python/sglang/srt/function_call/qwen25_detector.py new file mode 100644 index 000000000..40a65e5df --- /dev/null +++ b/python/sglang/srt/function_call/qwen25_detector.py @@ -0,0 +1,130 @@ +import json +import logging +import re +from typing import List + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer + +logger = logging.getLogger(__name__) + + +class Qwen25Detector(BaseFormatDetector): + """ + Detector for Qwen 2.5 and Qwen 3 model function call format. + + Format Structure: + ``` + \n{"name":"func1", "arguments":{...}}\n\n\n{"name":"func2", "arguments":{...}}\n + ``` + + Key Components: + - Tool Call Tags: `` and `` wrap each individual call + - Function Call Object: JSON object with "name" and "arguments" fields + + Reference: https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct?chat_template=default + """ + + def __init__(self): + """ + Initializes the detector with necessary state variables. + """ + super().__init__() + self.bot_token = "\n" + self.eot_token = "\n" + self.tool_call_separator = "\n" + self._normal_text_buffer = "" # Buffer for handling partial end tokens + + def has_tool_call(self, text: str) -> bool: + """Check if the text contains a Qwen 2.5 format tool call.""" + return self.bot_token in text + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """ + One-time parsing: Detects and parses tool calls in the provided text. + + :param text: The complete text to parse. + :param tools: List of available tools. + :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. + """ + idx = text.find(self.bot_token) + normal_text = text[:idx].strip() if idx != -1 else text + if self.bot_token not in text: + return StreamingParseResult(normal_text=normal_text, calls=[]) + + # Find all \n...\n blocks + pattern = rf"{re.escape(self.bot_token)}(.*?){re.escape(self.eot_token)}" + match_result_list = re.findall(pattern, text, re.DOTALL) + calls = [] + for match_result in match_result_list: + try: + parsed_call = json.loads(match_result.strip()) + calls.extend(self.parse_base_json(parsed_call, tools)) + except json.JSONDecodeError as e: + logger.warning( + f"Failed to parse JSON part: {match_result}, JSON parse error: {str(e)}" + ) + continue + return StreamingParseResult(normal_text=normal_text, calls=calls) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Streaming incremental parsing for Qwen 2.5 tool calls. + Uses base class implementation with buffering to handle partial end tokens. + """ + result = super().parse_streaming_increment(new_text, tools) + + # Handle partial end tokens that are streamed character by character + if result.normal_text: + self._normal_text_buffer += result.normal_text + + # Check if buffer contains complete end token (without leading newline) + end_token_without_newline = self.eot_token[1:] # "" + if end_token_without_newline in self._normal_text_buffer: + cleaned_text = self._normal_text_buffer.replace( + end_token_without_newline, "" + ) + self._normal_text_buffer = "" + result.normal_text = cleaned_text + else: + # Check if buffer might contain partial end token at the end + partial_match_len = self._ends_with_partial_token( + self._normal_text_buffer, end_token_without_newline + ) + + if partial_match_len: + # Keep potential partial match in buffer, return the rest + result.normal_text = self._normal_text_buffer[:-partial_match_len] + self._normal_text_buffer = self._normal_text_buffer[ + -partial_match_len: + ] + else: + # No partial match, return all buffered text + result.normal_text = self._normal_text_buffer + self._normal_text_buffer = "" + + return result + + def structure_info(self) -> _GetInfoFunc: + return lambda name: StructureInfo( + begin='\n{"name":"' + name + '", "arguments":', + end="}\n", + trigger="", + ) + + def build_ebnf(self, tools: List[Tool]): + return EBNFComposer.build_ebnf( + tools, + individual_call_start_token=self.bot_token.replace("\n", "\\n"), + individual_call_end_token=self.eot_token.replace("\n", "\\n"), + tool_call_separator="\\n", + function_format="json", + ) diff --git a/python/sglang/srt/function_call/qwen3_coder_detector.py b/python/sglang/srt/function_call/qwen3_coder_detector.py new file mode 100644 index 000000000..9bd3c7c24 --- /dev/null +++ b/python/sglang/srt/function_call/qwen3_coder_detector.py @@ -0,0 +1,362 @@ +import ast +import html +import json +import logging +import re +from typing import Any, Dict, List, Tuple + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer + +logger = logging.getLogger(__name__) + + +def _safe_val(raw: str) -> Any: + raw = html.unescape(raw.strip()) + try: + return json.loads(raw) + except Exception: + try: + return ast.literal_eval(raw) + except Exception: + return raw + + +class Qwen3CoderDetector(BaseFormatDetector): + """ + Detector for Qwen 3 models. + Assumes function call format: + + + + pwd && ls + + + + """ + + def __init__(self): + super().__init__() + self.tool_call_start_token: str = "" + self.tool_call_end_token: str = "" + self.tool_call_prefix: str = "(.*?)|(.*?)$", re.DOTALL + ) + self.tool_call_function_regex = re.compile( + r"|| bool: + return self.tool_call_start_token in text + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + normal, calls = self._extract(text, tools) + return StreamingParseResult(normal_text=normal, calls=calls) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + self._buf += new_text + normal = "" + calls: List[ToolCallItem] = [] + + # Build tool indices for validation + if not hasattr(self, "_tool_indices"): + self._tool_indices = self._get_tool_indices(tools) + + while True: + # If we're not in a tool call and don't see a start token, return normal text + if not self._in_tool_call and self.tool_call_start_token not in self._buf: + normal += self._buf + self._buf = "" + break + + # Look for tool call start + if not self._in_tool_call: + s = self._buf.find(self.tool_call_start_token) + if s == -1: + normal += self._buf + self._buf = "" + break + + normal += self._buf[:s] + self._buf = self._buf[s:] + + self._in_tool_call = True + self._function_name_sent = False + self._current_function_name = "" + self._current_parameters = {} + self._streamed_parameters = {} + + # Remove the start token + self._buf = self._buf[len(self.tool_call_start_token) :] + continue + + # We're in a tool call, try to parse function name if not sent yet + if not self._function_name_sent: + # Look for function name pattern: + function_match = re.search(r"]+)>", self._buf) + if function_match: + function_name = function_match.group(1).strip() + + # Validate function name + if function_name in self._tool_indices: + self._current_function_name = function_name + self._function_name_sent = True + + # Initialize tool call tracking + if self.current_tool_id == -1: + self.current_tool_id = 0 + + # Ensure tracking arrays are large enough + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + + # Store tool call info + self.prev_tool_call_arr[self.current_tool_id] = { + "name": function_name, + "arguments": {}, + } + + # Send tool name with empty parameters + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=function_name, + parameters="", + ) + ) + + # Remove the processed function declaration + self._buf = self._buf[function_match.end() :] + continue + else: + # Invalid function name, reset state + logger.warning(f"Invalid function name: {function_name}") + self._reset_streaming_state() + normal += self._buf + self._buf = "" + break + else: + # Function name not complete yet, wait for more text + break + + # Parse parameters incrementally + if self._function_name_sent: + # Process parameters and get any calls to emit + parameter_calls = self._parse_and_stream_parameters(self._buf) + calls.extend(parameter_calls) + + # Check if tool call is complete + if self.tool_call_end_token in self._buf: + end_pos = self._buf.find(self.tool_call_end_token) + + # Add closing brace to complete the JSON object + current_streamed = self.streamed_args_for_tool[self.current_tool_id] + if current_streamed: + # Count opening and closing braces to check if JSON is complete + open_braces = current_streamed.count("{") + close_braces = current_streamed.count("}") + if open_braces > close_braces: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=None, + parameters="}", + ) + ) + self.streamed_args_for_tool[self.current_tool_id] = ( + current_streamed + "}" + ) + + # Complete the tool call + self._buf = self._buf[end_pos + len(self.tool_call_end_token) :] + self._reset_streaming_state() + self.current_tool_id += 1 + continue + else: + # Tool call not complete yet, wait for more text + break + + return StreamingParseResult(normal_text=normal, calls=calls) + + def _parse_and_stream_parameters(self, text_to_parse: str) -> List[ToolCallItem]: + """ + Parse complete parameter blocks from text and return any tool call items to emit. + + This method: + 1. Finds all complete blocks + 2. Parses them into a dictionary + 3. Compares with current parameters and generates diff if needed + 4. Updates internal state + + Args: + text_to_parse: The text to search for parameter blocks + + Returns: + List of ToolCallItem objects to emit (may be empty) + """ + calls: List[ToolCallItem] = [] + + # Find all complete parameter patterns + param_matches = list( + re.finditer( + r"]+)>(.*?)", text_to_parse, re.DOTALL + ) + ) + + # Build new parameters dictionary + new_params = {} + for match in param_matches: + param_name = match.group(1).strip() + param_value = match.group(2) + new_params[param_name] = _safe_val(param_value) + + # Calculate parameter diff to stream with proper incremental JSON building + if new_params != self._current_parameters: + previous_args_json = self.streamed_args_for_tool[self.current_tool_id] + + # Build incremental JSON properly + if not self._current_parameters: + # First parameter(s) - start JSON object but don't close it yet + items = [] + for key, value in new_params.items(): + items.append( + f"{json.dumps(key, ensure_ascii=False)}: {json.dumps(value, ensure_ascii=False)}" + ) + json_fragment = "{" + ", ".join(items) + + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=None, + parameters=json_fragment, + ) + ) + self.streamed_args_for_tool[self.current_tool_id] = json_fragment + + else: + # Additional parameters - add them incrementally + new_keys = set(new_params.keys()) - set(self._current_parameters.keys()) + if new_keys: + # Build the continuation part (no closing brace yet) + continuation_parts = [] + for key in new_keys: + value = new_params[key] + continuation_parts.append( + f"{json.dumps(key, ensure_ascii=False)}: {json.dumps(value, ensure_ascii=False)}" + ) + + json_fragment = ", " + ", ".join(continuation_parts) + + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=None, + parameters=json_fragment, + ) + ) + self.streamed_args_for_tool[self.current_tool_id] = ( + previous_args_json + json_fragment + ) + + # Update current state + self._current_parameters = new_params + self.prev_tool_call_arr[self.current_tool_id]["arguments"] = new_params + + return calls + + def _reset_streaming_state(self): + """Reset streaming state for the next tool call""" + self._in_tool_call = False + self._function_name_sent = False + self._current_function_name = "" + self._current_parameters = {} + self._streamed_parameters = {} + self.current_tool_name_sent = False + + def _extract(self, text: str, tools: List[Tool]) -> Tuple[str, List[ToolCallItem]]: + normal_parts: List[str] = [] + calls: List[ToolCallItem] = [] + cursor = 0 + while True: + s = text.find(self.tool_call_start_token, cursor) + if s == -1: + normal_parts.append(text[cursor:]) + break + normal_parts.append(text[cursor:s]) + e = text.find(self.tool_call_end_token, s) + if e == -1: + normal_parts.append(text[s:]) + break + block = text[s : e + len(self.tool_call_end_token)] + cursor = e + len(self.tool_call_end_token) + calls.extend(self._parse_block(block, tools)) + return "".join(normal_parts), calls + + def _parse_block(self, block: str, tools: List[Tool]) -> List[ToolCallItem]: + res: List[ToolCallItem] = [] + for m in self.tool_call_function_regex.findall(block): + txt = m[0] if m[0] else m[1] + if ">" not in txt: + continue + idx = txt.index(">") + fname = txt[:idx].strip() + body = txt[idx + 1 :] + params: Dict[str, Any] = {} + for pm in self.tool_call_parameter_regex.findall(body): + ptxt = pm[0] if pm[0] else pm[1] + if ">" not in ptxt: + continue + pidx = ptxt.index(">") + pname = ptxt[:pidx].strip() + pval = ptxt[pidx + 1 :].lstrip("\n").rstrip("\n") + params[pname] = _safe_val(pval) + raw = {"name": fname, "arguments": params} + try: + # TODO: fix idx in function call, the index for a function + # call will always be -1 in parse_base_json + res.extend(self.parse_base_json(raw, tools)) + except Exception: + logger.warning("invalid tool call for %s dropped", fname) + return res + + def supports_structural_tag(self) -> bool: + return False + + def structure_info(self) -> _GetInfoFunc: + raise NotImplementedError + + def build_ebnf(self, tools: List[Tool]): + return EBNFComposer.build_ebnf( + tools, + individual_call_start_token=self.tool_call_start_token.replace("\n", "\\n"), + individual_call_end_token=self.tool_call_end_token.replace("\n", "\\n"), + tool_call_separator="\\n", + function_format="xml", + call_rule_fmt='"\\n" {arguments_rule} "\\n"', + key_value_rule_fmt='"\\n" {valrule} "\\n"', + key_value_separator='"\\n"', + ) diff --git a/python/sglang/srt/function_call/step3_detector.py b/python/sglang/srt/function_call/step3_detector.py new file mode 100644 index 000000000..b46f4544f --- /dev/null +++ b/python/sglang/srt/function_call/step3_detector.py @@ -0,0 +1,436 @@ +import ast +import json +import logging +import re +from typing import Any, Dict, List + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.function_call.ebnf_composer import EBNFComposer + +logger = logging.getLogger(__name__) + + +def get_argument_type(func_name: str, arg_key: str, defined_tools: List[Tool]) -> str: + """Get the expected type for a function argument from tool schema.""" + name2tool = {tool.function.name: tool for tool in defined_tools} + if func_name not in name2tool: + return None + tool = name2tool[func_name] + parameters = tool.function.parameters or {} + properties = parameters.get("properties", {}) + if arg_key not in properties: + return None + return properties[arg_key].get("type", None) + + +def parse_arguments(value: str) -> tuple[Any, bool]: + """Parse a string value to appropriate type. Returns (parsed_value, success).""" + try: + try: + parsed_value = json.loads(value) + except: + parsed_value = ast.literal_eval(value) + return parsed_value, True + except: + return value, False + + +class Step3Detector(BaseFormatDetector): + """ + Detector for Step3 model function call format. + + The Step3 format uses special Unicode tokens to delimit function calls + with steptml XML format for invocations. + + Format Structure: + ``` + <｜tool_calls_begin｜> + <｜tool_call_begin｜>function<｜tool_sep｜> + value1 + value2 + <｜tool_call_end｜> + <｜tool_calls_end｜> + ``` + """ + + def __init__(self): + super().__init__() + self.bot_token = "<｜tool_calls_begin｜>" + self.eot_token = "<｜tool_calls_end｜>" + self.tool_call_begin = "<｜tool_call_begin｜>" + self.tool_call_end = "<｜tool_call_end｜>" + self.tool_sep = "<｜tool_sep｜>" + + # Regex for parsing steptml invocations + self.invoke_regex = re.compile( + r'(.+?)', re.DOTALL + ) + self.param_regex = re.compile( + r'([^<]*)', re.DOTALL + ) + + # Streaming state variables + self._in_tool_block: bool = False + self._tool_block_finished: bool = False + self._current_function_name: str = "" + self._current_parameters: Dict[str, Any] = {} + self._in_tool_call: bool = False + self._function_name_sent: bool = False + + def has_tool_call(self, text: str) -> bool: + """Check if the text contains a Step3 format tool call.""" + return self.bot_token in text + + def _parse_steptml_invoke( + self, text: str, tools: List[Tool] = None + ) -> tuple[str, dict]: + """Parse steptml invoke format to extract function name and parameters.""" + invoke_match = self.invoke_regex.search(text) + if not invoke_match: + return None, {} + + func_name = invoke_match.group(1) + params_text = invoke_match.group(2) + + params = {} + for param_match in self.param_regex.finditer(params_text): + param_name = param_match.group(1) + param_value = param_match.group(2).strip() + + # If tools provided, use schema-aware parsing + if tools: + arg_type = get_argument_type(func_name, param_name, tools) + if arg_type and arg_type != "string": + parsed_value, _ = parse_arguments(param_value) + params[param_name] = parsed_value + else: + params[param_name] = param_value + else: + # Fallback to generic parsing if no tools provided + parsed_value, _ = parse_arguments(param_value) + params[param_name] = parsed_value + + return func_name, params + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """ + One-time parsing: Detects and parses tool calls in the provided text. + """ + if self.bot_token not in text: + return StreamingParseResult(normal_text=text, calls=[]) + + try: + pre_text, rest = text.split(self.bot_token, 1) + + # If no end token, return everything as normal text + if self.eot_token not in rest: + return StreamingParseResult(normal_text=text, calls=[]) + + tool_section, post_text = rest.split(self.eot_token, 1) + + # Find all individual tool calls using regex + calls = [] + tool_call_pattern = ( + f"{re.escape(self.tool_call_begin)}(.*?){re.escape(self.tool_call_end)}" + ) + + for match in re.finditer(tool_call_pattern, tool_section, re.DOTALL): + call_content = match.group(1) + + # Check if it's a function call + if self.tool_sep not in call_content: + continue + + type_part, invoke_part = call_content.split(self.tool_sep, 1) + if type_part.strip() != "function": + continue + + func_name, params = self._parse_steptml_invoke(invoke_part, tools) + if func_name: + # Use parse_base_json to create the ToolCallItem + action = {"name": func_name, "arguments": params} + calls.extend(self.parse_base_json(action, tools)) + + # Combine pre and post text + normal_text = pre_text + post_text + + return StreamingParseResult(normal_text=normal_text, calls=calls) + + except Exception as e: + logger.error(f"Error in detect_and_parse: {e}") + # Return the original text if parsing fails + return StreamingParseResult(normal_text=text) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Streaming incremental parsing for Step3 format. + """ + self._buffer += new_text + + # Build tool indices for validation + if not hasattr(self, "_tool_indices"): + self._tool_indices = self._get_tool_indices(tools) + + # If we've finished the tool block, everything is normal text + if self._tool_block_finished: + normal_text = self._buffer + self._buffer = "" + return StreamingParseResult(normal_text=normal_text) + + # Check if tool block hasn't started yet + if not self._in_tool_block: + if self.bot_token in self._buffer: + idx = self._buffer.find(self.bot_token) + normal_text = self._buffer[:idx] + self._buffer = self._buffer[idx + len(self.bot_token) :] + self._in_tool_block = True + return StreamingParseResult(normal_text=normal_text) + else: + # Check if we might have a partial bot_token + partial_len = self._ends_with_partial_token( + self._buffer, self.bot_token + ) + if partial_len: + return StreamingParseResult() # Wait for more text + else: + normal_text = self._buffer + self._buffer = "" + return StreamingParseResult(normal_text=normal_text) + + # We're inside the tool block + calls: List[ToolCallItem] = [] + + # Check if tool block is ending + if self.eot_token in self._buffer: + idx = self._buffer.find(self.eot_token) + + # If we're in the middle of a tool call, we need to handle it + if self._in_tool_call: + # The buffer before eot_token might contain the end of the current tool call + before_eot = self._buffer[:idx] + if self.tool_call_end in before_eot: + # Parse this final tool call + result = self._parse_partial_tool_call(tools) + calls.extend(result.calls) + else: + # Incomplete tool call - log warning + logger.warning("Tool block ended with incomplete tool call") + + remaining = self._buffer[idx + len(self.eot_token) :] + self._buffer = "" + self._tool_block_finished = True + + # Reset any partial tool call state + self._reset_streaming_state() + + return StreamingParseResult(normal_text=remaining, calls=calls) + + # Check if we're in a tool call or need to start one + if not self._in_tool_call: + if self.tool_call_begin in self._buffer: + idx = self._buffer.find(self.tool_call_begin) + # Remove any content before tool call begin (shouldn't happen but be safe) + self._buffer = self._buffer[idx + len(self.tool_call_begin) :] + self._in_tool_call = True + self._function_name_sent = False + self._current_function_name = "" + self._current_parameters = {} + # Fall through to parse the partial tool call + else: + # Wait for tool call to begin + return StreamingParseResult() + + # Parse partial tool call + if self._in_tool_call: + return self._parse_partial_tool_call(tools) + + return StreamingParseResult() + + def _parse_partial_tool_call(self, tools: List[Tool]) -> StreamingParseResult: + """Parse partial tool call for streaming scenarios.""" + calls = [] + + # Check if we have tool_sep (means we're past the type declaration) + if self.tool_sep not in self._buffer: + return StreamingParseResult(calls=calls) # Wait for more text + + type_part, invoke_part = self._buffer.split(self.tool_sep, 1) + if type_part.strip() != "function": + # Invalid tool type, skip this tool call + self._reset_streaming_state() + return StreamingParseResult(calls=calls) + + # Try to extract function name if not sent yet + if not self._function_name_sent: + name_match = re.search(r'', invoke_part) + if name_match: + func_name = name_match.group(1) + + # Validate function name + if func_name in self._tool_indices: + self._current_function_name = func_name + self._function_name_sent = True + + # Initialize tool tracking + if self.current_tool_id == -1: + self.current_tool_id = 0 + + # Ensure tracking arrays are large enough + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + + # Store tool call info + self.prev_tool_call_arr[self.current_tool_id] = { + "name": func_name, + "arguments": {}, + } + + # Send tool name with empty parameters + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=func_name, + parameters="", + ) + ) + else: + # Invalid function name + logger.warning(f"Invalid function name: {func_name}") + self._reset_streaming_state() + return StreamingParseResult(calls=calls) + else: + # Function name not complete yet + return StreamingParseResult(calls=calls) + + # Parse parameters incrementally + if self._function_name_sent: + # Extract all complete parameters + new_params = {} + for param_match in self.param_regex.finditer(invoke_part): + param_name = param_match.group(1) + param_value = param_match.group(2).strip() + + # Use schema-aware parsing + arg_type = get_argument_type( + self._current_function_name, param_name, tools + ) + if arg_type and arg_type != "string": + parsed_value, _ = parse_arguments(param_value) + new_params[param_name] = parsed_value + else: + new_params[param_name] = param_value + + # Check if we have new parameters to stream + if new_params != self._current_parameters: + # Build the JSON content without the closing brace for streaming + if not self._current_parameters: + # First parameters - send opening brace and content + params_content = json.dumps(new_params, ensure_ascii=False) + if len(params_content) > 2: # More than just "{}" + # Send everything except the closing brace + diff = params_content[:-1] + else: + diff = "{" + else: + # Subsequent parameters - calculate the incremental diff + old_json = json.dumps(self._current_parameters, ensure_ascii=False) + new_json = json.dumps(new_params, ensure_ascii=False) + + # Remove closing braces for comparison + old_without_brace = old_json[:-1] + new_without_brace = new_json[:-1] + + # The new content should extend the old content + if new_without_brace.startswith(old_without_brace): + diff = new_without_brace[len(old_without_brace) :] + else: + # Parameters changed in unexpected way - shouldn't happen in normal streaming + diff = "" + + if diff: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + parameters=diff, + ) + ) + self.streamed_args_for_tool[self.current_tool_id] += diff + + # Update current state + self._current_parameters = new_params + self.prev_tool_call_arr[self.current_tool_id]["arguments"] = new_params + + # Check if tool call is complete + if self.tool_call_end in self._buffer: + # Send closing brace if we've sent any parameters + if self.streamed_args_for_tool[self.current_tool_id]: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + parameters="}", + ) + ) + self.streamed_args_for_tool[self.current_tool_id] += "}" + + # Find the end position + end_idx = self._buffer.find(self.tool_call_end) + # Remove the processed tool call from buffer + self._buffer = self._buffer[end_idx + len(self.tool_call_end) :] + + # Reset state for next tool call + self._reset_streaming_state() + self.current_tool_id += 1 + + return StreamingParseResult(calls=calls) + + def _reset_streaming_state(self): + """Reset streaming state for the next tool call""" + self._in_tool_call = False + self._function_name_sent = False + self._current_function_name = "" + self._current_parameters = {} + + def supports_structural_tag(self) -> bool: + """Return True if this detector supports structural tag format.""" + return False + + def structure_info(self) -> _GetInfoFunc: + raise NotImplementedError() + + def build_ebnf(self, tools: List[Tool]) -> str: + """ + Build EBNF grammar for Step3 tool call format. + """ + # Custom call rule for steptml format + call_rule_fmt = ( + '"function" "<｜tool_sep｜>" "" ' + '{arguments_rule} ""' + ) + + # Custom key-value rule for steptml parameters + key_value_rule_fmt = ( + '"" {valrule} ""' + ) + + return EBNFComposer.build_ebnf( + tools, + sequence_start_token=self.bot_token, + sequence_end_token=self.eot_token, + individual_call_start_token=self.tool_call_begin, + individual_call_end_token=self.tool_call_end, + tool_call_separator="", + function_format="xml", + call_rule_fmt=call_rule_fmt, + key_value_rule_fmt=key_value_rule_fmt, + key_value_separator="", + ) diff --git a/python/sglang/srt/function_call/utils.py b/python/sglang/srt/function_call/utils.py new file mode 100644 index 000000000..c4da456f3 --- /dev/null +++ b/python/sglang/srt/function_call/utils.py @@ -0,0 +1,52 @@ +import json +from json import JSONDecodeError, JSONDecoder +from typing import Any, Tuple + +import partial_json_parser +from partial_json_parser.core.options import Allow + + +def _find_common_prefix(s1: str, s2: str) -> str: + prefix = "" + min_length = min(len(s1), len(s2)) + for i in range(0, min_length): + if s1[i] == s2[i]: + prefix += s1[i] + else: + break + return prefix + + +def _partial_json_loads(input_str: str, flags: Allow) -> Tuple[Any, int]: + """ + Parse incomplete or partial JSON strings commonly encountered during streaming. + + Args: + input_str (str): The potentially incomplete JSON string to parse. + flags (Allow): Bitwise flags controlling what types of partial data are allowed. + Common flags include: + - Allow.STR: Allow partial strings (e.g., '"hello wo' -> 'hello wo') + - Allow.OBJ: Allow partial objects (e.g., '{"key":' -> {'key': None}) + - Allow.ARR: Allow partial arrays (e.g., '[1, 2,' -> [1, 2]) + - Allow.ALL: Allow all types of partial data + + Returns: + Tuple[Any, int]: A tuple containing: + - parsed_object: The Python object parsed from the JSON + - consumed_length: Number of characters consumed from input_str + """ + try: + return (partial_json_parser.loads(input_str, flags), len(input_str)) + except JSONDecodeError as e: + if "Extra data" in e.msg: + dec = JSONDecoder() + return dec.raw_decode(input_str) + raise + + +def _is_complete_json(input_str: str) -> bool: + try: + json.loads(input_str) + return True + except JSONDecodeError: + return False diff --git a/python/sglang/srt/hf_transformers_utils.py b/python/sglang/srt/hf_transformers_utils.py new file mode 100644 index 000000000..d7dcf8904 --- /dev/null +++ b/python/sglang/srt/hf_transformers_utils.py @@ -0,0 +1,442 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Utilities for Huggingface Transformers.""" + +import contextlib +import json +import os +import warnings +from pathlib import Path +from typing import Any, Dict, Optional, Type, Union + +import torch +from huggingface_hub import snapshot_download +from transformers import ( + AutoConfig, + AutoProcessor, + AutoTokenizer, + GenerationConfig, + PretrainedConfig, + PreTrainedTokenizer, + PreTrainedTokenizerBase, + PreTrainedTokenizerFast, +) +from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES + +from sglang.srt.configs import ( + ChatGLMConfig, + DbrxConfig, + DeepseekVL2Config, + ExaoneConfig, + KimiVLConfig, + LongcatFlashConfig, + MultiModalityConfig, + Qwen3NextConfig, + Step3VLConfig, +) +from sglang.srt.configs.internvl import InternVLChatConfig +from sglang.srt.connector import create_remote_connector +from sglang.srt.utils import is_remote_url, logger, lru_cache_frozenset + +_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { + ChatGLMConfig.model_type: ChatGLMConfig, + DbrxConfig.model_type: DbrxConfig, + ExaoneConfig.model_type: ExaoneConfig, + DeepseekVL2Config.model_type: DeepseekVL2Config, + MultiModalityConfig.model_type: MultiModalityConfig, + KimiVLConfig.model_type: KimiVLConfig, + InternVLChatConfig.model_type: InternVLChatConfig, + Step3VLConfig.model_type: Step3VLConfig, + LongcatFlashConfig.model_type: LongcatFlashConfig, + Qwen3NextConfig.model_type: Qwen3NextConfig, +} + +for name, cls in _CONFIG_REGISTRY.items(): + with contextlib.suppress(ValueError): + AutoConfig.register(name, cls) + + +def download_from_hf( + model_path: str, + allow_patterns: Optional[Union[str, list]] = None, +): + if os.path.exists(model_path): + return model_path + + if not allow_patterns: + allow_patterns = ["*.json", "*.bin", "*.model"] + + return snapshot_download(model_path, allow_patterns=allow_patterns) + + +def get_hf_text_config(config: PretrainedConfig): + """Get the "sub" config relevant to llm for multi modal models. + No op for pure text models. + """ + if config.architectures is not None: + class_name = config.architectures[0] + if class_name.startswith("Llava") and class_name.endswith("ForCausalLM"): + # We support non-hf version of llava models, so we do not want to + # read the wrong values from the unused default text_config. + # NOTE(HandH1998): We set `torch_dtype` of config to `torch.float16` for the weights, as + # `torch.float16` is default used for image features in `python/sglang/srt/models/llava.py`. + setattr(config, "torch_dtype", torch.float16) + return config + + if hasattr(config, "text_config"): + # The code operates under the assumption that text_config should have + # `num_attention_heads` (among others). Assert here to fail early + # if transformers config doesn't align with this assumption. + assert hasattr(config.text_config, "num_attention_heads") + return config.text_config + if hasattr(config, "language_config"): + return config.language_config + if hasattr(config, "thinker_config"): + # qwen2.5 omni + thinker_config = config.thinker_config + if hasattr(thinker_config, "text_config"): + setattr( + thinker_config.text_config, + "torch_dtype", + getattr(thinker_config, "torch_dtype", None), + ) + return thinker_config.text_config + return thinker_config + else: + return config + + +@lru_cache_frozenset(maxsize=32) +def get_config( + model: str, + trust_remote_code: bool, + revision: Optional[str] = None, + model_override_args: Optional[dict] = None, + **kwargs, +): + is_gguf = check_gguf_file(model) + if is_gguf: + kwargs["gguf_file"] = model + model = Path(model).parent + + if is_remote_url(model): + # BaseConnector implements __del__() to clean up the local dir. + # Since config files need to exist all the time, so we DO NOT use + # with statement to avoid closing the client. + client = create_remote_connector(model) + client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"]) + model = client.get_local_dir() + + config = AutoConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + if ( + config.architectures is not None + and config.architectures[0] == "Phi4MMForCausalLM" + ): + # Phi4MMForCausalLM uses a hard-coded vision_config. See: + # https://github.com/vllm-project/vllm/blob/6071e989df1531b59ef35568f83f7351afb0b51e/vllm/model_executor/models/phi4mm.py#L71 + # We set it here to support cases where num_attention_heads is not divisible by the TP size. + from transformers import SiglipVisionConfig + + vision_config = { + "hidden_size": 1152, + "image_size": 448, + "intermediate_size": 4304, + "model_type": "siglip_vision_model", + "num_attention_heads": 16, + "num_hidden_layers": 26, # Model is originally 27-layer, we only need the first 26 layers for feature extraction. + "patch_size": 14, + } + config.vision_config = SiglipVisionConfig(**vision_config) + text_config = get_hf_text_config(config=config) + + if isinstance(model, str) and text_config is not None: + for key, val in text_config.__dict__.items(): + if not hasattr(config, key) and getattr(text_config, key, None) is not None: + setattr(config, key, val) + + if config.model_type in _CONFIG_REGISTRY: + config_class = _CONFIG_REGISTRY[config.model_type] + config = config_class.from_pretrained(model, revision=revision) + # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`. + setattr(config, "_name_or_path", model) + + if isinstance(model, str) and config.model_type == "internvl_chat": + for key, val in config.llm_config.__dict__.items(): + if not hasattr(config, key): + setattr(config, key, val) + + if config.model_type == "multi_modality": + config.update({"architectures": ["MultiModalityCausalLM"]}) + + if model_override_args: + config.update(model_override_args) + + # Special architecture mapping check for GGUF models + if is_gguf: + if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: + raise RuntimeError(f"Can't get gguf config for {config.model_type}.") + model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type] + config.update({"architectures": [model_type]}) + + return config + + +@lru_cache_frozenset(maxsize=32) +def get_generation_config( + model: str, + trust_remote_code: bool, + revision: Optional[str] = None, + **kwargs, +): + try: + return GenerationConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + except OSError as e: + return None + + +# Qwen-1M related +def get_sparse_attention_config( + model: str, + sparse_attention_config_filename: str = "sparse_attention_config.json", +) -> Dict[str, Any]: + is_local = os.path.isdir(model) + if not is_local: + # Download the config files. + model = download_from_hf(model, allow_patterns=["*.json"]) + + config_file = os.path.join(model, sparse_attention_config_filename) + if not os.path.exists(config_file): + return {} + + # Load the sparse attention config. + with open(config_file) as f: + config = json.load(f) + return config + + +# Models don't use the same configuration key for determining the maximum +# context length. Store them here so we can sanely check them. +# NOTE: The ordering here is important. Some models have two of these and we +# have a preference for which value gets used. +CONTEXT_LENGTH_KEYS = [ + "max_sequence_length", + "seq_length", + "max_seq_len", + "model_max_length", + "max_position_embeddings", +] + + +def get_context_length(config): + """Get the context length of a model from a huggingface model configs.""" + text_config = config + rope_scaling = getattr(text_config, "rope_scaling", None) + if rope_scaling: + rope_scaling_factor = rope_scaling.get("factor", 1) + if "original_max_position_embeddings" in rope_scaling: + rope_scaling_factor = 1 + if rope_scaling.get("rope_type", None) == "llama3": + rope_scaling_factor = 1 + else: + rope_scaling_factor = 1 + + for key in CONTEXT_LENGTH_KEYS: + val = getattr(text_config, key, None) + if val is not None: + return int(rope_scaling_factor * val) + return 2048 + + +# A fast LLaMA tokenizer with the pre-processed `tokenizer.json` file. +_FAST_LLAMA_TOKENIZER = "hf-internal-testing/llama-tokenizer" + + +def get_tokenizer( + tokenizer_name: str, + *args, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + tokenizer_revision: Optional[str] = None, + **kwargs, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + """Gets a tokenizer for the given model name via Huggingface.""" + if tokenizer_name.endswith(".json"): + from sglang.srt.tokenizer.tiktoken_tokenizer import TiktokenTokenizer + + return TiktokenTokenizer(tokenizer_name) + + if tokenizer_mode == "slow": + if kwargs.get("use_fast", False): + raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.") + kwargs["use_fast"] = False + + # TODO(Xinyuan): Remove this once we have a proper tokenizer for Devstral + if tokenizer_name == "mistralai/Devstral-Small-2505": + tokenizer_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503" + + is_gguf = check_gguf_file(tokenizer_name) + if is_gguf: + kwargs["gguf_file"] = tokenizer_name + tokenizer_name = Path(tokenizer_name).parent + + if is_remote_url(tokenizer_name): + # BaseConnector implements __del__() to clean up the local dir. + # Since config files need to exist all the time, so we DO NOT use + # with statement to avoid closing the client. + client = create_remote_connector(tokenizer_name) + client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"]) + tokenizer_name = client.get_local_dir() + + try: + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + tokenizer_revision=tokenizer_revision, + clean_up_tokenization_spaces=False, + **kwargs, + ) + except TypeError as e: + # The LLaMA tokenizer causes a protobuf error in some environments. + err_msg = ( + "Failed to load the tokenizer. If you are using a LLaMA V1 model " + f"consider using '{_FAST_LLAMA_TOKENIZER}' instead of the " + "original tokenizer." + ) + raise RuntimeError(err_msg) from e + except ValueError as e: + # If the error pertains to the tokenizer class not existing or not + # currently being imported, suggest using the --trust-remote-code flag. + if not trust_remote_code and ( + "does not exist or is not currently imported." in str(e) + or "requires you to execute the tokenizer file" in str(e) + ): + err_msg = ( + "Failed to load the tokenizer. If the tokenizer is a custom " + "tokenizer not yet available in the HuggingFace transformers " + "library, consider setting `trust_remote_code=True` in LLM " + "or using the `--trust-remote-code` flag in the CLI." + ) + raise RuntimeError(err_msg) from e + else: + raise e + + if not isinstance(tokenizer, PreTrainedTokenizerFast): + warnings.warn( + "Using a slow tokenizer. This might cause a significant " + "slowdown. Consider using a fast tokenizer instead." + ) + + attach_additional_stop_token_ids(tokenizer) + return tokenizer + + +# Some models doesn't have an available processor, e.g.: InternVL +def get_tokenizer_from_processor(processor): + if isinstance(processor, PreTrainedTokenizerBase): + return processor + return processor.tokenizer + + +def get_processor( + tokenizer_name: str, + *args, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + tokenizer_revision: Optional[str] = None, + use_fast: Optional[bool] = True, + **kwargs, +): + # pop 'revision' from kwargs if present. + revision = kwargs.pop("revision", tokenizer_revision) + + config = AutoConfig.from_pretrained( + tokenizer_name, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + + # fix: for Qwen2-VL model, inject default 'size' if not provided. + if config.model_type in {"qwen2_vl"}: + if "size" not in kwargs: + kwargs["size"] = {"shortest_edge": 3136, "longest_edge": 1003520} + + if config.model_type not in {"llava", "clip"}: + kwargs["use_fast"] = use_fast + try: + if "InternVL3_5" in tokenizer_name: + processor = AutoTokenizer.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + else: + processor = AutoProcessor.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + + except ValueError as e: + error_message = str(e) + if "does not have a slow version" in error_message: + logger.info( + f"Processor {tokenizer_name} does not have a slow version. Automatically use fast version" + ) + kwargs["use_fast"] = True + processor = AutoProcessor.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + else: + raise e + tokenizer = get_tokenizer_from_processor(processor) + + attach_additional_stop_token_ids(tokenizer) + return processor + + +def attach_additional_stop_token_ids(tokenizer): + # Special handling for stop token <|eom_id|> generated by llama 3 tool use. + if "<|eom_id|>" in tokenizer.get_added_vocab(): + tokenizer.additional_stop_token_ids = set( + [tokenizer.get_added_vocab()["<|eom_id|>"]] + ) + else: + tokenizer.additional_stop_token_ids = None + + +def check_gguf_file(model: Union[str, os.PathLike]) -> bool: + """Check if the file is a GGUF model.""" + model = Path(model) + if not model.is_file(): + return False + elif model.suffix == ".gguf": + return True + + with open(model, "rb") as f: + header = f.read(4) + return header == b"GGUF" diff --git a/python/sglang/srt/host_shared_memory.py b/python/sglang/srt/host_shared_memory.py new file mode 100644 index 000000000..c599527f9 --- /dev/null +++ b/python/sglang/srt/host_shared_memory.py @@ -0,0 +1,83 @@ +import logging +import os +from dataclasses import dataclass +from multiprocessing import shared_memory +from pathlib import Path +from typing import List, Optional + +import numpy as np +import torch + +from sglang.srt.distributed.naive_distributed import get_naive_distributed +from sglang.srt.utils import check_cuda_result + +logger = logging.getLogger(__name__) + + +class HostSharedMemoryManager: + def __init__(self, base_name: str): + self._base_name = Path(base_name) + self._operation_index = 0 + self._records: List[_Record] = [] + + def malloc(self, *, shape, dtype): + meta_tensor = torch.empty(size=shape, dtype=dtype, device="meta") + raw = self._malloc_raw(num_bytes=meta_tensor.nbytes) + return raw.view(dtype).view(*shape) + + def _malloc_raw(self, *, num_bytes: int) -> torch.Tensor: + import cuda.bindings.runtime as cuda_rt + + self._operation_index += 1 + shm_name = f"{self._base_name}_op{self._operation_index}" + + # TODO handle dispose + if get_naive_distributed().get_rank() == 0: + shm = shared_memory.SharedMemory(name=shm_name, create=True, size=num_bytes) + + get_naive_distributed().barrier() + + if get_naive_distributed().get_rank() != 0: + shm = shared_memory.SharedMemory(name=shm_name) + + np_array = np.ndarray((num_bytes,), dtype=np.uint8, buffer=shm.buf) + tensor = torch.from_numpy(np_array) + + check_cuda_result( + cuda_rt.cudaHostRegister( + tensor.data_ptr(), num_bytes, cuda_rt.cudaHostRegisterPortable + ) + ) + + get_naive_distributed().barrier() + + self._records.append( + _Record( + shm=shm, + np_array=np_array, + tensor=tensor, + ) + ) + return tensor + + +@dataclass +class _Record: + shm: shared_memory.SharedMemory + np_array: np.ndarray + tensor: torch.Tensor + + +# Can have multi instances if needed +_instance: Optional[HostSharedMemoryManager] = None + + +def get_host_shared_memory_manager(): + assert _instance is not None + return _instance + + +def set_host_shared_memory_manager(instance: HostSharedMemoryManager): + global _instance + assert _instance is None + _instance = instance diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py new file mode 100644 index 000000000..3d973393e --- /dev/null +++ b/python/sglang/srt/layers/activation.py @@ -0,0 +1,272 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Fused operators for activation layers.""" + +import logging +import math +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import PretrainedConfig + +from sglang.srt.custom_op import CustomOp +from sglang.srt.distributed import ( + divide, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.utils import ( + cpu_has_amx_support, + is_cpu, + is_cuda, + is_hip, + is_npu, + is_xpu, + set_weight_attrs, +) +from sglang.utils import resolve_obj_by_qualname + +_is_cuda = is_cuda() +_is_npu = is_npu() +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() +_is_hip = is_hip() +_is_xpu = is_xpu() + +if _is_cuda or _is_xpu: + from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul +elif _is_hip: + from sgl_kernel import gelu_and_mul, gelu_quick, gelu_tanh_and_mul, silu_and_mul + +if is_npu(): + import torch_npu + +logger = logging.getLogger(__name__) + + +class SiluAndMul(CustomOp): + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + silu_and_mul(x, out) + return out + + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if _is_cpu_amx_available: + out = torch.ops.sgl_kernel.silu_and_mul_cpu(x) + return out + else: + return self.forward_native(x) + + def forward_npu(self, x: torch.Tensor) -> torch.Tensor: + out = torch_npu.npu_swiglu(x) + return out + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + silu_and_mul(x, out) + return out + + +class GeluAndMul(CustomOp): + def __init__(self, approximate="tanh"): + super().__init__() + self.approximate = approximate + + def _forward_impl(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + if self.approximate == "tanh": + gelu_tanh_and_mul(x, out) + elif self.approximate == "none": + gelu_and_mul(x, out) + else: + raise RuntimeError("GeluAndMul only support tanh or none") + return out + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:] + + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if _is_cpu_amx_available and self.approximate == "tanh": + return torch.ops.sgl_kernel.gelu_tanh_and_mul_cpu(x) + elif _is_cpu_amx_available and self.approximate == "none": + return torch.ops.sgl_kernel.gelu_and_mul_cpu(x) + else: + return self.forward_native(x) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + return self._forward_impl(x) + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self._forward_impl(x) + + def forward_npu(self, x: torch.Tensor) -> torch.Tensor: + y_npu, gelu_npu = torch_npu.npu_geglu( + x, + dim=-1, + approximate=1 if self.approximate == "tanh" else 0, + activate_left=True, + ) + return y_npu + + +class NewGELU(CustomOp): + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + c = math.sqrt(2.0 / math.pi) + return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * torch.pow(x, 3.0)))) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + # TODO: Implement the CUDA kernel for NewGELU in sgl-kernel + return self.forward_native(x) + + +class ReLU2(nn.Module): + """ + Applies the squared Rectified Linear Unit function. + y = max(0, x)^2 + """ + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.relu(x) + return x * x + + +class QuickGELU(CustomOp): + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.sigmoid(1.702 * x) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + def forward_hip(self, x: torch.Tensor) -> torch.Tensor: + out = torch.empty(x.shape, dtype=x.dtype, device=x.device) + gelu_quick(x, out) + return out + + def forward_npu(self, x: torch.Tensor) -> torch.Tensor: + return torch_npu.npu_fast_gelu(x) + + +class ScaledActivation(nn.Module): + """An activation function with post-scale parameters. + + This is used for some quantization methods like AWQ. + """ + + def __init__( + self, + act_module: nn.Module, + intermediate_size: int, + input_is_parallel: bool = True, + params_dtype: Optional[torch.dtype] = None, + ): + super().__init__() + self.act = act_module + self.input_is_parallel = input_is_parallel + if input_is_parallel: + tp_size = get_tensor_model_parallel_world_size() + intermediate_size_per_partition = divide(intermediate_size, tp_size) + else: + intermediate_size_per_partition = intermediate_size + if params_dtype is None: + params_dtype = torch.get_default_dtype() + self.scales = nn.Parameter( + torch.empty(intermediate_size_per_partition, dtype=params_dtype) + ) + set_weight_attrs(self.scales, {"weight_loader": self.weight_loader}) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(x) / self.scales + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + param_data = param.data + if self.input_is_parallel: + tp_rank = get_tensor_model_parallel_rank() + shard_size = param_data.shape[0] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow(0, start_idx, shard_size) + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +_ACTIVATION_REGISTRY = { + "gelu": nn.GELU(), + "gelu_pytorch_tanh": nn.GELU(approximate="tanh"), + "gelu_new": NewGELU(), + "relu2": ReLU2(), +} + + +def get_act_fn( + act_fn_name: str, + quant_config: Optional[QuantizationConfig] = None, + intermediate_size: Optional[int] = None, + input_is_parallel: bool = True, + params_dtype: Optional[torch.dtype] = None, +) -> nn.Module: + """Get an activation function by name.""" + act_fn_name = act_fn_name.lower() + if act_fn_name not in _ACTIVATION_REGISTRY: + raise ValueError(f"Activation function {act_fn_name!r} is not supported.") + + act_fn = _ACTIVATION_REGISTRY[act_fn_name] + if quant_config is not None and act_fn_name in quant_config.get_scaled_act_names(): + if intermediate_size is None: + raise ValueError( + "intermediate_size must be specified for scaled " + "activation functions." + ) + return ScaledActivation( + act_fn, intermediate_size, input_is_parallel, params_dtype + ) + return act_fn + + +def get_cross_encoder_activation_function(config: PretrainedConfig): + if ( + hasattr(config, "sbert_ce_default_activation_function") + and config.sbert_ce_default_activation_function is not None + ): + + function_name = config.sbert_ce_default_activation_function + assert function_name.startswith("torch.nn.modules."), ( + "Loading of activation functions is restricted to " + "torch.nn.modules for security reasons" + ) + return resolve_obj_by_qualname(function_name)() + else: + # adapt bge-reranker + return nn.Identity() + + +if not ( + _is_cuda or _is_npu or (_is_cpu and _is_cpu_amx_available) or _is_hip or _is_xpu +): + logger.info( + "sgl-kernel is not available on Non-NV, Non-AMD platforms or Non-AMX CPUs. Fallback to other kernel libraries." + ) + from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul diff --git a/python/sglang/srt/layers/amx_utils.py b/python/sglang/srt/layers/amx_utils.py new file mode 100644 index 000000000..df2a05ba5 --- /dev/null +++ b/python/sglang/srt/layers/amx_utils.py @@ -0,0 +1,86 @@ +import logging + +import torch + +from sglang.srt.utils import cpu_has_amx_support + +logger = logging.getLogger(__name__) + + +def amx_process_weight_after_loading(weight): + if weight.device != torch.device("cpu"): + return weight + if not cpu_has_amx_support(): + return weight + + return torch.ops.sgl_kernel.convert_weight_packed(weight) + + +# TODO: currently gemm kernel has the below requirements: +# OC % TILE_N == 0, where TILE_N = 16 +# IC % TILE_K == 0, where TILE_K = 32 +def dim_is_supported(weight): + TILE_N = 16 + TILE_K = 32 + ndim = weight.ndim + OC = weight.size(1) if ndim == 3 else weight.size(0) + IC = weight.size(2) if ndim == 3 else weight.size(1) + return OC % TILE_N == 0 and IC % TILE_K == 0 + + +def _amx_process_weight_after_loading( + module, weight_names, transpose_dims=None +) -> None: + # Pack weight for get better performance on CPU + devices = {getattr(module, weight_name).device for weight_name in weight_names} + assert len(devices) == 1, f"Expects all weights to be on the same device" + device = devices.pop() + + if transpose_dims: + assert len(weight_names) == len( + transpose_dims + ), "len(weight_names) should be equal to len(transpose_dims)" + + for i, weight_name in enumerate(weight_names): + weight_tensor = getattr(module, weight_name) + + if transpose_dims and transpose_dims[i]: + weight_tensor = weight_tensor.transpose(*transpose_dims[i]) + + # We don't pack weight or use intel amx backend if any weight of this module has unsupported dim. + if not dim_is_supported(weight_tensor): + logger.warning( + f"Unsupported dimension for prepacking for weight '{weight_name}' with shape {weight_tensor.shape} in {module}. " + f"The derived (OC, IC) dimensions must be divisible by (16, 32). " + ) + module.use_intel_amx_backend = False + return + + packed_weight = torch.nn.Parameter( + amx_process_weight_after_loading(weight_tensor), + requires_grad=False, + ) + packed_weight.__dict__ = weight_tensor.__dict__ + setattr(module, weight_name, packed_weight) + + module.use_intel_amx_backend = ( + device == torch.device("cpu") and cpu_has_amx_support() + ) + + if ( + module.use_intel_amx_backend + and hasattr(module, "bias") + and module.bias is not None + ): + module.bias = torch.nn.Parameter(module.bias.data.float(), requires_grad=False) + + +class PackWeightMethod: + def __init__(self, weight_names, transpose_dims=None): + self.weight_names = weight_names + self.transpose_dims = transpose_dims + + def process_weights_after_loading(self, module) -> None: + _amx_process_weight_after_loading( + module, self.weight_names, self.transpose_dims + ) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py new file mode 100644 index 000000000..188d772c7 --- /dev/null +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -0,0 +1,1179 @@ +from __future__ import annotations + +""" +end to end attention solution with aiter kernels +""" + +import math +import os +from dataclasses import dataclass +from enum import Enum, auto +from functools import partial +from typing import TYPE_CHECKING, List, Optional, Union + +import torch +import triton +import triton.language as tl + +from sglang.global_config import global_config +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton +from sglang.srt.layers.dp_attention import ( + get_attention_tp_size, + is_dp_attention_enabled, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpecInfo + +try: + from aiter import ( + flash_attn_varlen_func, + mha_batch_prefill_func, + paged_attention_ragged, + ) + from aiter.mla import mla_decode_fwd, mla_prefill_fwd +except ImportError: + print( + "aiter is AMD specific kernel library. Please make sure aiter is installed on your AMD device." + ) + +from sglang.srt.configs.model_config import AttentionArch + + +class WrapperDispatch(Enum): + SLIDING_WINDOW = auto() + CROSS_ATTENTION = auto() + + +@dataclass +class ForwardMetadata: + kv_indptr: torch.Tensor + kv_indices: torch.Tensor + qo_indptr: torch.Tensor + kv_last_page_len: torch.Tensor + max_q_len: int + max_kv_len: Optional[int] + + +global_workspace_buffer = None + +_AITER_PARTITION_SIZE_ROCM = 256 + + +class AiterAttnBackend(AttentionBackend): + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + kv_indptr_buf: Optional[torch.Tensor] = None, + ): + super().__init__() + # Lazy import to avoid the initialization of cuda context + from sglang.srt.layers.attention.triton_ops.extend_attention import ( + extend_attention_fwd, + ) + + self.extend_attention_fwd = torch.compiler.disable(extend_attention_fwd) + + self.device = model_runner.device + self.is_multimodal = model_runner.model_config.is_multimodal + self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens + self.speculative_num_steps = model_runner.server_args.speculative_num_steps + self.num_head = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.head_dim = model_runner.model_config.head_dim + self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1] + self.num_kv_head = model_runner.model_config.get_num_kv_heads( + get_attention_tp_size() + ) + self.kv_cache_dtype = model_runner.kv_cache_dtype + + self.req_to_token = model_runner.req_to_token_pool.req_to_token + + self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA + + # Parse constants + self.max_context_len = model_runner.model_config.context_len + self.skip_prefill = skip_prefill + + max_bs = model_runner.req_to_token_pool.size + + if kv_indptr_buf is None: + self.kv_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + else: + self.kv_indptr = kv_indptr_buf + + self.kv_last_page_len = torch.ones( + (max_bs,), dtype=torch.int32, device=model_runner.device + ) + self.qo_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + + # Create prefill indices updater + if not skip_prefill: + self.indices_updater_prefill = AiterIndicesUpdaterPrefill( + model_runner, self + ) + if self.use_mla: + self.mla_indices_updater_prefill = AiterMlaIndicesUpdaterPrefill( + model_runner, self + ) + + # aiter kernel related initialization + self.max_num_partitions = ( + self.max_context_len + _AITER_PARTITION_SIZE_ROCM - 1 + ) // _AITER_PARTITION_SIZE_ROCM + + nbyes_per_qo_elem = torch.finfo(torch.float32).bits // 8 + + if not self.use_mla: + self.workspace_buffer = torch.empty( + (max_bs * self.num_head * self.max_num_partitions * self.head_dim) + * nbyes_per_qo_elem + + 2 * (max_bs * self.num_head * self.max_num_partitions) * 4, + dtype=torch.uint8, + device=self.device, + ) + + self.scale = float(1.0 / (self.head_dim**0.5)) + self.k_scale = self.v_scale = torch.tensor([1.0], dtype=torch.float32).to( + self.device + ) + + self.logits_soft_cap = 0.0 + + self.forward_metadata: ForwardMetadata = None + + if self.use_mla: + self.qo_indptr_ = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + + self.enable_dp_attention = is_dp_attention_enabled() + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init auxiliary variables for triton attention backend.""" + + bs = forward_batch.batch_size + kv_indptr = self.kv_indptr + spec_info = forward_batch.spec_info + qo_indptr = None + kv_last_page_len = None + max_q_len = None + + if forward_batch.forward_mode.is_decode_or_idle(): + if spec_info is None: + kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + forward_batch.seq_lens_sum, dtype=torch.int32, device=self.device + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + else: + kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices + bs = kv_indptr.shape[0] - 1 + + if self.use_mla: + qo_indptr = self.qo_indptr_[: bs + 1] + qo_indptr[1 : bs + 1] = torch.cumsum(self.kv_last_page_len[:bs], dim=0) + kv_last_page_len = self.kv_last_page_len[:bs] + max_q_len = 1 + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + None, + ) + + elif forward_batch.forward_mode.is_draft_extend(): + if self.use_mla: + kv_indices, kv_indptr, qo_indptr, custom_mask = ( + spec_info.generate_attn_arg_prefill( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + self.req_to_token, + ) + ) + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + # self.mla_indices_updater_prefill.kv_last_page_len, + self.kv_last_page_len[:bs], + max(forward_batch.extend_seq_lens_cpu), + forward_batch.seq_lens_cpu.max().item(), + ) + else: + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + prefix_lens=None, + encoder_lens=forward_batch.encoder_lens, + spec_info=forward_batch.spec_info, + ) + self.forward_metadata = ForwardMetadata( + self.indices_updater_prefill.kv_indptr, + self.indices_updater_prefill.kv_indices, + None, + None, + self.indices_updater_prefill.max_q_len, + self.indices_updater_prefill.max_kv_len, + ) + elif forward_batch.forward_mode.is_target_verify(): + if self.use_mla: + draft_num = spec_info.draft_token_num + kv_lens = forward_batch.seq_lens + draft_num + kv_lens_sum = forward_batch.seq_lens_sum + draft_num * bs + device = forward_batch.seq_lens.device + + qo_indptr = torch.arange( + 0, + (1 + bs) * draft_num, + step=draft_num, + dtype=torch.int32, + device=device, + ) + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + kv_lens_sum, + dtype=torch.int32, + device=device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + kv_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + # self.mla_indices_updater_prefill.kv_last_page_len, + self.kv_last_page_len[:bs], + draft_num, + None, + ) + else: + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + prefix_lens=None, + encoder_lens=forward_batch.encoder_lens, + spec_info=forward_batch.spec_info, + ) + self.forward_metadata = ForwardMetadata( + self.indices_updater_prefill.kv_indptr, + self.indices_updater_prefill.kv_indices, + None, + None, + self.indices_updater_prefill.max_q_len, + self.indices_updater_prefill.max_kv_len, + ) + else: + prefix_lens = forward_batch.extend_prefix_lens + + if self.is_multimodal: + extend_no_prefix = False + else: + extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu) + if self.use_mla: + self.mla_indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + forward_batch.extend_seq_lens, + forward_batch.extend_seq_lens.max().item(), + forward_batch.seq_lens.max().item(), + spec_info=None, + ) + + kv_indices = self.mla_indices_updater_prefill.kv_indices + + self.forward_metadata = ForwardMetadata( + self.mla_indices_updater_prefill.kv_indptr, + kv_indices, + self.mla_indices_updater_prefill.qo_indptr, + self.kv_last_page_len[:bs], + self.mla_indices_updater_prefill.max_q_len, + self.mla_indices_updater_prefill.max_kv_len, + ) + else: + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + prefix_lens, + encoder_lens=forward_batch.encoder_lens, + spec_info=None, + ) + self.forward_metadata = ForwardMetadata( + self.indices_updater_prefill.kv_indptr, + self.indices_updater_prefill.kv_indices, + None, + None, + self.indices_updater_prefill.max_q_len, + self.indices_updater_prefill.max_kv_len, + ) + + def init_cuda_graph_state( + self, + max_bs: int, + max_num_tokens: int, + kv_indices_buf: Optional[torch.Tensor] = None, + ): + self.cuda_graph_kv_last_page_len = torch.ones(max_bs, dtype=torch.int) + if kv_indices_buf is None: + self.cuda_graph_kv_indices = torch.zeros( + (max_bs * self.max_context_len), + dtype=torch.int32, + device=self.device, + ) + else: + self.cuda_graph_kv_indices = kv_indices_buf + + if not self.skip_prefill: + self.cuda_graph_custom_mask = torch.zeros( + (max_num_tokens * self.max_context_len), + dtype=torch.uint8, + device=self.device, + ) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + ): + if forward_mode.is_decode_or_idle(): + qo_indptr = None + kv_last_page_len = None + max_q_len = None + + if spec_info is None: + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + else: + kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices + + if self.use_mla: + qo_indptr = self.qo_indptr_[: bs + 1] + qo_indptr[1 : bs + 1] = torch.cumsum( + self.cuda_graph_kv_last_page_len[:bs], dim=0 + ) + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = 1 + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + None, + ) + + elif forward_mode.is_target_verify(): + if self.use_mla: + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = self.num_draft_tokens + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + None, + ) + else: + seq_lens_sum = seq_lens.sum().item() + self.indices_updater_prefill.update( + req_pool_indices, + seq_lens, + seq_lens_sum, + prefix_lens=None, + encoder_lens=encoder_lens, + spec_info=spec_info, + ) + self.forward_metadata = ForwardMetadata( + self.indices_updater_prefill.kv_indptr, + self.indices_updater_prefill.kv_indices, + None, + None, + self.indices_updater_prefill.max_q_len, + self.indices_updater_prefill.max_kv_len, + ) + elif forward_mode.is_draft_extend(): + num_tokens_per_bs = self.speculative_num_steps + 1 + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + bs * num_tokens_per_bs + 1, + step=num_tokens_per_bs, + dtype=torch.int32, + device=self.device, + ) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = num_tokens_per_bs + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + None, + ) + else: + raise ValueError(f"Invalid mode: {forward_mode=}") + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + seq_lens_cpu: Optional[torch.Tensor], + ): + if forward_mode.is_decode_or_idle(): + kv_indptr = self.kv_indptr + kv_indices = self.cuda_graph_kv_indices + if spec_info is None: + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens[:bs], dim=0) + kv_indptr = kv_indptr[: bs + 1] + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices[:bs], + seq_lens[:bs], + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + else: + kv_indptr[: spec_info.kv_indptr.shape[0]] = spec_info.kv_indptr + kv_indices[: spec_info.kv_indices.shape[0]] = spec_info.kv_indices + + elif forward_mode.is_target_verify(): + bs = len(req_pool_indices) + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_lens = seq_lens + self.num_draft_tokens + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + kv_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + elif forward_mode.is_draft_extend(): + seq_lens = seq_lens[:bs] + accept_lens = spec_info.accept_length[:bs] + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[1 : bs + 1] = torch.cumsum(accept_lens, dim=0) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + else: + raise ValueError("Invalid forward mode") + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + + self.logits_soft_cap = layer.logit_cap + + if k is not None: + assert v is not None + if save_kv_cache: + if self.use_mla: + forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) + else: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, layer.k_scale, layer.v_scale + ) + + if self.use_mla: + max_q_len = self.forward_metadata.max_q_len + max_kv_len = self.forward_metadata.max_kv_len + kv_indptr = self.forward_metadata.kv_indptr + kv_indices = self.forward_metadata.kv_indices + qo_indptr = self.forward_metadata.qo_indptr + K_Buffer = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + V_Buffer = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id) + kv_lora_rank = V_Buffer.shape[-1] + qk_rope_head_dim = K_Buffer.shape[-1] - kv_lora_rank + qk_nope_head_dim = k.shape[-1] - qk_rope_head_dim + assert len(q.shape) == 3 + assert len(k.shape) == 3 + assert len(v.shape) == 3 + + if forward_batch.forward_mode.is_extend(): + if kv_indices.shape[0] == 0: + o = flash_attn_varlen_func( + q, + k, + v, + qo_indptr, + qo_indptr, + max_q_len, + max_q_len, + softmax_scale=layer.scaling, + causal=True, + ) + return o + elif layer.qk_head_dim != (kv_lora_rank + qk_rope_head_dim): + K_Buffer = torch.index_select(K_Buffer, 0, kv_indices) + kvc, k_pe = torch.split( + K_Buffer, [kv_lora_rank, qk_rope_head_dim], dim=-1 + ) + kvprefix = layer.kv_b_proj(kvc.contiguous())[0] + + kvprefix = kvprefix.view( + -1, layer.tp_k_head_num, qk_nope_head_dim + layer.v_head_dim + ) + k_prefix, v_prefix = torch.split( + kvprefix, [qk_nope_head_dim, layer.v_head_dim], dim=-1 + ) + k_prefix = torch.cat( + [ + k_prefix, + torch.broadcast_to( + k_pe, + (k_pe.shape[0], layer.tp_k_head_num, k_pe.shape[2]), + ), + ], + dim=-1, + ) + assert ( + forward_batch.extend_prefix_lens.shape + == forward_batch.extend_seq_lens.shape + ) + + k = k_prefix + v = v_prefix + + o = flash_attn_varlen_func( + q, + k, + v, + qo_indptr, + kv_indptr, + max_q_len, + max_kv_len, + softmax_scale=layer.scaling, + causal=True, + ) + return o + + else: + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty( + (q.shape[0], layer.tp_q_head_num * layer.v_head_dim) + ) + else: + o = torch.empty_like(q) + + mla_prefill_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + K_Buffer.view(-1, 1, 1, layer.qk_head_dim), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + qo_indptr, + kv_indptr, + kv_indices, + self.forward_metadata.kv_last_page_len, + self.forward_metadata.max_q_len, + layer.scaling, + layer.logit_cap, + ) + K_Buffer = K_Buffer.view(-1, layer.tp_k_head_num, layer.qk_head_dim) + return o + elif forward_batch.forward_mode.is_target_verify(): + o = q.new_empty((q.shape[0], layer.tp_q_head_num, layer.v_head_dim)) + mla_decode_fwd( + q, + K_Buffer.view(-1, 1, 1, layer.qk_head_dim), + o, + self.forward_metadata.qo_indptr, + self.forward_metadata.kv_indptr, + self.forward_metadata.kv_indices, + self.forward_metadata.kv_last_page_len, + self.forward_metadata.max_q_len, + layer.scaling, + layer.logit_cap, + ) + K_Buffer = K_Buffer.view(-1, 1, layer.qk_head_dim) + return o + elif forward_batch.forward_mode.is_draft_extend(): + o = q.new_empty((q.shape[0], layer.tp_q_head_num, layer.v_head_dim)) + causal = True + sliding_window_size = -1 + kv_indptr = self.forward_metadata.kv_indptr + kv_indices = self.forward_metadata.kv_indices + mla_prefill_fwd( + q, + K_Buffer.view(-1, 1, 1, layer.qk_head_dim), + o, + self.forward_metadata.qo_indptr, + self.forward_metadata.kv_indptr, + self.forward_metadata.kv_indices, + self.forward_metadata.kv_last_page_len, + self.forward_metadata.max_q_len, + layer.scaling, + layer.logit_cap, + ) + K_Buffer = K_Buffer.view(-1, 1, layer.qk_head_dim) + return o + # self.extend_attention_fwd( + # q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + # k.contiguous(), + # v.contiguous(), + # o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + # forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + # forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + # self.forward_metadata.qo_indptr, + # kv_indptr, + # kv_indices, + # None, + # causal, + # None, + # self.forward_metadata.max_q_len, + # layer.scaling, + # layer.logit_cap, + # sliding_window_size, + # ) + # return o + else: + raise ValueError( + f"Invalid forward mode for MLA prefill: {forward_batch.forward_mode=}" + ) + else: + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer( + layer.layer_id + ) + + bs0 = forward_batch.batch_size + 1 + + o = mha_batch_prefill_func( + q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache, + v_cache, + self.qo_indptr[:bs0], + self.forward_metadata.kv_indptr[:bs0], + self.forward_metadata.kv_indices, + self.forward_metadata.max_q_len, + self.forward_metadata.max_kv_len, + causal=True, + logits_soft_cap=self.logits_soft_cap, + alibi_slopes=None, + return_lse=False, + return_attn_probs=False, + ) + + return o.view(-1, layer.tp_q_head_num * layer.head_dim) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + + q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim) + + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + + if self.use_mla: + k_buffer = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + mla_decode_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + k_buffer.view(-1, 1, 1, layer.qk_head_dim), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + self.forward_metadata.qo_indptr, + self.forward_metadata.kv_indptr, + self.forward_metadata.kv_indices, + self.forward_metadata.kv_last_page_len, + self.forward_metadata.max_q_len, + layer.scaling, + layer.logit_cap, + ) + k_buffer = k_buffer.view(-1, 1, layer.qk_head_dim) + else: + self.logits_soft_cap = layer.logit_cap + paged_attention_ragged( + o.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + self.workspace_buffer, + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id).view( + -1, 1, layer.tp_k_head_num, layer.qk_head_dim + ), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id).view( + -1, 1, layer.tp_v_head_num, layer.v_head_dim + ), + self.scale, + self.forward_metadata.kv_indptr, + self.forward_metadata.kv_indices, + self.kv_last_page_len, + 1, + self.max_num_partitions, + None, + "auto", + "NHD", + self.logits_soft_cap, + self.k_scale, + self.v_scale, + None, + _AITER_PARTITION_SIZE_ROCM, + ) + + return o + + +class AiterIndicesUpdaterPrefill: + def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend): + # Parse Constants + self.num_qo_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.num_kv_heads = model_runner.model_config.get_num_kv_heads( + get_attention_tp_size() + ) + self.head_dim = model_runner.model_config.head_dim + self.data_type = model_runner.kv_cache_dtype + self.q_data_type = model_runner.dtype + self.sliding_window_size = model_runner.sliding_window_size + self.attn_backend = attn_backend + + # Buffers and wrappers + self.kv_indptr = attn_backend.kv_indptr + self.kv_last_page_len = attn_backend.kv_last_page_len + self.qo_indptr = attn_backend.qo_indptr + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.update = self.update_single_wrapper + + self.kv_indices = None + self.max_q_len = 0 + self.max_kv_len = 0 + + def update( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + prefix_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[SpecInfo], + ): + # Keep the signature for type checking. It will be assigned during runtime. + raise NotImplementedError() + + def update_single_wrapper( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + prefix_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[SpecInfo], + ): + + kv_start_idx = None + kv_indptr = self.kv_indptr + qo_indptr = self.qo_indptr + paged_kernel_lens = seq_lens + paged_kernel_lens_sum = seq_lens_sum + + bs = len(req_pool_indices) + if spec_info is None: + # Normal extend + kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + + # (TODO: Kk) WA - CI test_moe_eval_accuracy_large.py + # mha_batch_prefill reads 128 data to do computatoin + # if real data is not long enough then original padding value 0 is used + # but the 0 location will be made nan (noqa) in cuda graph capture mode + # this will cause the output tensor value becomes nan + # WA is to assure that last index of pool not changed + kv_indices = torch.empty( + paged_kernel_lens_sum + 256, + dtype=torch.int32, + device=req_pool_indices.device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + paged_kernel_lens, + kv_indptr, + kv_start_idx, + kv_indices, + self.req_to_token.shape[1], + ) + + token_num = kv_indptr[-1] + kv_indices[token_num:] = kv_indices[0] + + self.max_kv_len = torch.max(paged_kernel_lens).item() + + extend_lens = seq_lens - prefix_lens + self.max_q_len = torch.max(extend_lens).item() + + qo_indptr[1 : bs + 1] = torch.cumsum(extend_lens, dim=0) + qo_indptr = qo_indptr[: bs + 1] + custom_mask = None + else: + kv_indices, kv_indptr, qo_indptr, custom_mask = ( + spec_info.generate_attn_arg_prefill( + req_pool_indices, + paged_kernel_lens, + paged_kernel_lens_sum, + self.req_to_token, + ) + ) + + self.kv_indices = kv_indices + + +class AiterMlaIndicesUpdaterPrefill: + def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend): + # Parse Constants + self.attn_backend = attn_backend + + # Buffers and wrappers + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.update = self.update_single_wrapper + + self.kv_indptr = None + self.kv_indices = None + self.qo_indptr = None + self.kv_last_page_len = None + self.max_q_len = 0 + self.max_kv_len = 0 + + def update( + self, + req_pool_indices: torch.Tensor, + kv_lens: torch.Tensor, + kv_lens_sum: int, + extend_lens: torch.Tensor, + max_q_len: int, + max_kv_len: int, + spec_info: Optional[SpecInfo], + ): + # Keep the signature for type checking. It will be assigned during runtime. + raise NotImplementedError() + + def update_single_wrapper( + self, + req_pool_indices: torch.Tensor, + kv_lens: torch.Tensor, + kv_lens_sum: int, + extend_lens: torch.Tensor, + max_q_len: int, + max_kv_len: int, + spec_info: Optional[SpecInfo], + ): + bs = len(req_pool_indices) + + kv_indptr = self.attn_backend.kv_indptr + + if spec_info is None: + # Normal extend + kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + kv_lens_sum, + dtype=torch.int32, + device=req_pool_indices.device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + kv_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + + qo_indptr = self.attn_backend.qo_indptr + qo_indptr[1 : bs + 1] = torch.cumsum(extend_lens, dim=0) + qo_indptr = qo_indptr[: bs + 1] + else: + kv_indices, kv_indptr, qo_indptr, custom_mask = ( + spec_info.generate_attn_arg_prefill( + req_pool_indices, + kv_lens, + kv_lens_sum, + self.req_to_token, + ) + ) + + self.kv_indptr = kv_indptr + self.kv_indices = kv_indices + self.qo_indptr = qo_indptr + self.max_q_len = max_q_len + self.max_kv_len = max_kv_len + + +class AiterMultiStepDraftBackend: + """ + Wrap multiple triton attention backends as one for multiple consecutive + draft decoding steps. + """ + + def __init__( + self, + model_runner: ModelRunner, + topk: int, + speculative_num_steps: int, + ): + from sglang.srt.speculative.eagle_utils import generate_draft_decode_kv_indices + + self.topk = topk + self.speculative_num_steps = speculative_num_steps + self.generate_draft_decode_kv_indices = generate_draft_decode_kv_indices + max_bs = model_runner.req_to_token_pool.size * self.topk + self.kv_indptr = torch.zeros( + ( + self.speculative_num_steps, + max_bs + 1, + ), + dtype=torch.int32, + device=model_runner.device, + ) + self.attn_backends = [] + for i in range(self.speculative_num_steps): + self.attn_backends.append( + AiterAttnBackend( + model_runner, + skip_prefill=True, + kv_indptr_buf=self.kv_indptr[i], + ) + ) + self.max_context_len = self.attn_backends[0].max_context_len + self.num_head = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.device = model_runner.device + # Cached variables for generate_draft_decode_kv_indices + self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1] + self.page_size = model_runner.server_args.page_size + assert self.page_size == 1, "Page size must be 1" + + def common_template( + self, forward_batch: ForwardBatch, kv_indices_buffer: torch.Tensor, call_fn: int + ): + num_seqs = forward_batch.batch_size + bs = self.topk * num_seqs + seq_lens_sum = forward_batch.seq_lens_sum + + self.generate_draft_decode_kv_indices[ + (self.speculative_num_steps, num_seqs, self.topk) + ]( + forward_batch.req_pool_indices, + forward_batch.req_to_token_pool.req_to_token, + forward_batch.seq_lens, + kv_indices_buffer, + self.kv_indptr, + forward_batch.positions, + self.pool_len, + kv_indices_buffer.shape[1], + self.kv_indptr.shape[1], + triton.next_power_of_2(num_seqs), + triton.next_power_of_2(self.speculative_num_steps), + triton.next_power_of_2(bs), + self.page_size, + ) + + for i in range(self.speculative_num_steps): + forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1] + forward_batch.spec_info.kv_indices = kv_indices_buffer[i][ + : seq_lens_sum * self.topk + bs * (i + 1) + ] + call_fn(i, forward_batch) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + kv_indices = torch.empty( + ( + self.speculative_num_steps, + forward_batch.batch_size * self.topk * self.max_context_len, + ), + dtype=torch.int32, + device=self.device, + ) + + def call_fn(i, forward_batch): + forward_batch.spec_info.kv_indptr = ( + forward_batch.spec_info.kv_indptr.clone() + ) + forward_batch.spec_info.kv_indices = ( + forward_batch.spec_info.kv_indices.clone() + ) + self.attn_backends[i].init_forward_metadata(forward_batch) + + self.common_template(forward_batch, kv_indices, call_fn) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.cuda_graph_kv_indices = torch.zeros( + (self.speculative_num_steps, max_num_tokens * self.max_context_len), + dtype=torch.int32, + device=self.device, + ) + for i in range(self.speculative_num_steps): + self.attn_backends[i].init_cuda_graph_state( + max_bs, max_num_tokens, kv_indices_buf=self.cuda_graph_kv_indices[i] + ) + + def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_capture_cuda_graph( + forward_batch.batch_size, + forward_batch.batch_size * self.topk, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + ) + + self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn) + + def init_forward_metadata_replay_cuda_graph( + self, forward_batch: ForwardBatch, bs: int + ): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_replay_cuda_graph( + bs, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + seq_lens_sum=-1, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + seq_lens_cpu=None, + ) + + self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn) diff --git a/python/sglang/srt/layers/attention/ascend_backend.py b/python/sglang/srt/layers/attention/ascend_backend.py new file mode 100644 index 000000000..6d5ed0a5c --- /dev/null +++ b/python/sglang/srt/layers/attention/ascend_backend.py @@ -0,0 +1,577 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional + +import torch +import torch_npu +from torch.nn.functional import scaled_dot_product_attention + +from sglang.srt.configs.model_config import AttentionArch +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.layers.radix_attention import AttentionType +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.utils import get_bool_env_var + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + +import os + +import numpy as np + + +@dataclass +class ForwardMetadata: + + # calculated map for kv positions [bs * maxseqlen] + block_tables: Optional[torch.Tensor] = None + + # seq len inputs + extend_seq_lens_cpu_int: Optional[torch.Tensor] = None + seq_lens_cpu_int: Optional[torch.Tensor] = None + seq_lens_cpu_list: Optional[List[int]] = None + seq_lens_list_cumsum: Optional[List[int]] = None + + +class AscendAttnBackend(AttentionBackend): + + def gen_attention_mask(self, max_seq_len: int, dtype=torch.float16): + mask_flag = torch.tril( + torch.ones((max_seq_len, max_seq_len), dtype=torch.bool) + ).view(max_seq_len, max_seq_len) + mask_flag = ~mask_flag + if dtype == torch.float16: + mask_value = torch.finfo(torch.float32).min + else: + mask_value = 1 + self.mask = ( + torch.masked_fill( + torch.zeros(size=(max_seq_len, max_seq_len)), mask_flag, mask_value + ) + .to(dtype) + .to(self.device) + ) + self.mask_len = max_seq_len + + def __init__(self, model_runner: ModelRunner): + super().__init__() + self.forward_metadata = None + self.device = model_runner.device + self.page_size = model_runner.page_size + self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA + if self.use_mla: + self.kv_lora_rank = model_runner.model_config.kv_lora_rank + self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim + self.native_attn = TorchNativeAttnBackend(model_runner) + self.graph_metadata = {} + self.max_context_len = model_runner.model_config.context_len + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.graph_mode = False + self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") + if not self.use_fia: + self.gen_attention_mask(128, model_runner.dtype) + mask_length = 2048 + self.fia_mask = ~torch.tril( + torch.ones( + (mask_length, mask_length), + dtype=torch.bool, + device=model_runner.device, + ) + ) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init the metadata for a forward pass.""" + tp_size = get_attention_tp_size() + self.forward_metadata = ForwardMetadata() + + self.forward_metadata.block_tables = ( + forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : forward_batch.seq_lens.max() + ][:, :: self.page_size] + // self.page_size + ) + if forward_batch.extend_seq_lens is not None: + self.forward_metadata.extend_seq_lens_cpu_int = ( + forward_batch.extend_seq_lens.cpu().int() + ) + self.forward_metadata.seq_lens_cpu_int = forward_batch.seq_lens_cpu.int() + + seq_lens_list_cumsum = np.cumsum(forward_batch.extend_seq_lens_cpu) + if forward_batch.is_extend_in_batch: + seq_lens_list_cumsum[-1] = ( + (seq_lens_list_cumsum[-1] - 1) // tp_size + 1 + ) * tp_size + self.forward_metadata.seq_lens_list_cumsum = seq_lens_list_cumsum + + self.graph_mode = False + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.graph_metadata = { + "block_tables": torch.empty( + (max_bs, self.max_context_len // self.page_size), + dtype=torch.int32, + device=self.device, + ), + } + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + metadata = ForwardMetadata() + + metadata.block_tables = self.graph_metadata["block_tables"][:bs, :] + metadata.seq_lens_cpu_list = seq_lens.cpu().int().tolist() + + self.graph_metadata[bs] = metadata + self.forward_metadata = metadata + + self.graph_mode = True + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + ): + metadata = self.graph_metadata[bs] + max_len = seq_lens_cpu[:bs].max().item() + max_seq_pages = (max_len + self.page_size - 1) // self.page_size + + metadata.block_tables[:bs, :max_seq_pages].copy_( + self.req_to_token[req_pool_indices[:bs], :max_len][:, :: self.page_size] + // self.page_size + ) + metadata.block_tables[:bs, max_seq_pages:].fill_(0) + metadata.block_tables[bs:, :].fill_(0) + + self.forward_metadata = metadata + + self.graph_mode = True + + def get_cuda_graph_seq_len_fill_value(self): + return 0 + + def forward_extend( + self, + q, + k, + v, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + ): + if not self.use_mla: + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + + k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + v_cache = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id) + + if self.use_fia: + """FIA will support multi-bs in the later version of CANN""" + q = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) + attn_output = torch.empty( + (q.size(0), layer.tp_q_head_num, layer.v_head_dim), + device=q.device, + dtype=q.dtype, + ) + q_len_offset = 0 + for q_len in forward_batch.extend_seq_lens_cpu: + attn_output[q_len_offset : q_len_offset + q_len] = ( + torch.ops.npu.npu_fused_infer_attention_score( + q[None, q_len_offset : q_len_offset + q_len], + k[None, q_len_offset : q_len_offset + q_len], + v[None, q_len_offset : q_len_offset + q_len], + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="BSND", # todo, TND not supports q_heads!=k_heads + atten_mask=self.fia_mask.unsqueeze(0), + sparse_mode=3, + scale=layer.scaling, + next_tokens=0, + )[0] + ) + q_len_offset += q_len + attn_output = attn_output.view( + -1, layer.tp_q_head_num * layer.v_head_dim + ) + + else: + if layer.qk_head_dim <= 128: + query = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim) + attn_output = torch.empty( + (query.shape[0], layer.tp_q_head_num * layer.v_head_dim), + dtype=query.dtype, + device=query.device, + ) + + torch_npu._npu_flash_attention_qlens( + query=query, + key_cache=k_cache, + value_cache=v_cache, + mask=self.mask, + block_table=self.forward_metadata.block_tables, + seq_len=self.forward_metadata.extend_seq_lens_cpu_int, + context_lens=self.forward_metadata.seq_lens_cpu_int, + scale_value=layer.scaling, + num_heads=layer.tp_q_head_num, + num_kv_heads=layer.tp_k_head_num, + out=attn_output, + ) + else: + if layer.qk_head_dim != layer.v_head_dim: + attn_output = q.new_empty( + (q.shape[0], layer.tp_q_head_num * layer.v_head_dim) + ) + else: + attn_output = torch.empty_like(q) + + use_gqa = layer.tp_q_head_num != layer.tp_k_head_num + + q_ = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) + o_ = attn_output.view(-1, layer.tp_q_head_num, layer.v_head_dim) + + causal = True + if ( + layer.is_cross_attention + or layer.attn_type == AttentionType.ENCODER_ONLY + ): + causal = False + + self.native_attn._run_sdpa_forward_extend( + q_, + o_, + k_cache.view(-1, layer.tp_k_head_num, layer.qk_head_dim), + v_cache.view(-1, layer.tp_v_head_num, layer.v_head_dim), + forward_batch.req_to_token_pool.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.extend_prefix_lens, + forward_batch.extend_seq_lens, + scaling=layer.scaling, + enable_gqa=use_gqa, + causal=causal, + ) + else: + assert ( + layer.qk_head_dim != layer.v_head_dim + ), "FIA only supports qk_head_dim != v_head_dim" + q_nope, q_rope = q.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1) + k_nope, k_rope = k.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1) + + attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score( + q_nope, + k_nope, + v, + query_rope=q_rope, + key_rope=k_rope, + num_heads=layer.tp_q_head_num, + input_layout="TND", + atten_mask=self.fia_mask, + sparse_mode=3, + actual_seq_lengths=self.forward_metadata.seq_lens_list_cumsum, + actual_seq_lengths_kv=self.forward_metadata.seq_lens_list_cumsum, + scale=layer.scaling, + next_tokens=0, + ) + + return attn_output + + def forward_decode_graph( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + ): + if save_kv_cache: + if self.use_mla: + k = k.view(-1, layer.tp_k_head_num, self.kv_lora_rank) + k_rope = k_rope.view(-1, layer.tp_k_head_num, self.qk_rope_head_dim) + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, k_rope + ) + else: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + + if not self.use_mla: + k_cache = forward_batch.token_to_kv_pool.get_key_buffer( + layer.layer_id + ).view(-1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim) + v_cache = forward_batch.token_to_kv_pool.get_value_buffer( + layer.layer_id + ).view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim) + query = q.reshape(-1, 1, layer.tp_q_head_num * layer.qk_head_dim) + if self.forward_metadata.seq_lens_cpu_int is None: + actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list + else: + actual_seq_len_kv = ( + self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist() + ) + num_tokens = query.shape[0] + workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace( + query, + k_cache, + v_cache, + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="BSH", + scale=layer.scaling, + actual_seq_lengths_kv=actual_seq_len_kv, + ) + output = torch.empty( + (num_tokens, 1, layer.tp_q_head_num * layer.v_head_dim), + dtype=q.dtype, + device=q.device, + ) + softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device) + torch_npu.npu_fused_infer_attention_score.out( + query, + k_cache, + v_cache, + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="BSH", + scale=layer.scaling, + actual_seq_lengths_kv=actual_seq_len_kv, + workspace=workspace, + out=[output, softmax_lse], + ) + return output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim) + else: + c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) + k_rope_cache = k_rope.view( + -1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim + ) + c_kv_cache = c_kv.view( + -1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank + ) + + q_nope = q.view(-1, layer.tp_q_head_num, 1, self.kv_lora_rank).contiguous() + q_rope = q_rope.view(-1, layer.tp_q_head_num, 1, self.qk_rope_head_dim) + if self.forward_metadata.seq_lens_cpu_int is None: + actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list + else: + actual_seq_len_kv = ( + self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist() + ) + + workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace( + q_nope, + c_kv_cache, + c_kv_cache, + query_rope=q_rope, + key_rope=k_rope_cache, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + input_layout="BNSD", + scale=layer.scaling, + actual_seq_lengths_kv=actual_seq_len_kv, + antiquant_mode=0, + antiquant_scale=None, + sparse_mode=0, + ) + output = torch.zeros_like(q_nope, dtype=q.dtype, device=q.device) + softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device) + + torch_npu.npu_fused_infer_attention_score.out( + q_nope, + c_kv_cache, + c_kv_cache, + query_rope=q_rope, + key_rope=k_rope_cache, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + input_layout="BNSD", + scale=layer.scaling, + actual_seq_lengths_kv=actual_seq_len_kv, + antiquant_mode=0, + antiquant_scale=None, + sparse_mode=0, + workspace=workspace, + out=[output, softmax_lse], + ) + return output.view(-1, layer.tp_q_head_num * self.kv_lora_rank) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + ): + if self.graph_mode: + return self.forward_decode_graph( + q, + k, + v, + layer, + forward_batch, + save_kv_cache, + q_rope=q_rope, + k_rope=k_rope, + ) + + if not self.use_mla: + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + num_tokens = q.shape[0] + k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + v_cache = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id) + if self.use_fia: + attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score( + q.view( + forward_batch.batch_size, + -1, + layer.tp_q_head_num, + layer.qk_head_dim, + ), + k_cache.view( + -1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim + ), + v_cache.view( + -1, self.page_size, layer.tp_v_head_num * layer.qk_head_dim + ), + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="BSND", + atten_mask=None, + block_size=self.page_size, + block_table=self.forward_metadata.block_tables, + actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_int, + scale=layer.scaling, + ) + else: + query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) + num_tokens = query.shape[0] + attn_output = torch.empty( + (num_tokens, layer.tp_q_head_num, layer.v_head_dim), + dtype=query.dtype, + device=query.device, + ) + + torch_npu._npu_paged_attention( + query=query, + key_cache=k_cache, + value_cache=v_cache, + num_heads=layer.tp_q_head_num, + num_kv_heads=layer.tp_k_head_num, + scale_value=layer.scaling, + block_table=self.forward_metadata.block_tables, + context_lens=self.forward_metadata.seq_lens_cpu_int, + out=attn_output, + ) + return attn_output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim) + else: + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, k_rope + ) + num_tokens = q.shape[0] + kv_c = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + k_pe = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id) + + if self.use_fia and (layer.tp_q_head_num // layer.tp_k_head_num) >= 8: + """layer.tp_q_head_num // layer.tp_k_head_num < 8 will support in the later version of CANN""" + kv_c = kv_c.view( + -1, self.page_size, layer.tp_k_head_num * self.kv_lora_rank + ) + k_pe = k_pe.view( + -1, self.page_size, layer.tp_k_head_num * self.qk_rope_head_dim + ) + q = q.view( + forward_batch.batch_size, -1, layer.tp_q_head_num, self.kv_lora_rank + ) + q_rope = q_rope.view( + forward_batch.batch_size, + -1, + layer.tp_q_head_num, + self.qk_rope_head_dim, + ) + attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score( + q, + kv_c, + kv_c, + query_rope=q_rope, + key_rope=k_pe, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="BSND", + atten_mask=None, + sparse_mode=0, + scale=layer.scaling, + antiquant_mode=0, + antiquant_scale=None, + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_int, + ) + else: + assert ( + self.graph_mode == False + ) # _npu_paged_attention_mla not support graph mode + q = torch.cat([q, q_rope], dim=-1) + query = q.view(-1, layer.tp_q_head_num, layer.head_dim) + kv_c_and_k_pe_cache = torch.cat([kv_c, k_pe], dim=-1) + kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view( + -1, + self.page_size, + layer.tp_k_head_num, + self.kv_lora_rank + self.qk_rope_head_dim, + ) + attn_output = torch.empty( + [num_tokens, layer.tp_q_head_num, self.kv_lora_rank], + dtype=q.dtype, + device=q.device, + ) + torch_npu._npu_paged_attention_mla( + query=query, + key_cache=kv_c_and_k_pe_cache, + num_kv_heads=layer.tp_k_head_num, + num_heads=layer.tp_q_head_num, + scale_value=layer.scaling, + block_table=self.forward_metadata.block_tables, + context_lens=self.forward_metadata.seq_lens_cpu_int, + mla_vheadsize=self.kv_lora_rank, + out=attn_output, + ) + return attn_output.view(num_tokens, layer.tp_q_head_num * self.kv_lora_rank) diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py new file mode 100644 index 000000000..3025d0b11 --- /dev/null +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Optional, Union + +import torch + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode + from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput + + +class AttentionBackend(ABC): + """The base class of attention backends""" + + @abstractmethod + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init the metadata for a forward pass.""" + raise NotImplementedError() + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + """Init the global shared states for cuda graph.""" + raise NotImplementedError() + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + """Init the metadata for a forward pass for capturing a cuda graph.""" + raise NotImplementedError() + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + ): + """Init the metadata for a forward pass for replaying a cuda graph.""" + raise NotImplementedError() + + def get_cuda_graph_seq_len_fill_value(self): + """Get the fill value for padded seq lens. Typically, it is 0 or 1.""" + raise NotImplementedError() + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + """Run forward on an attention layer.""" + if forward_batch.forward_mode.is_idle(): + return q.new_empty(q.shape[0], layer.tp_q_head_num * layer.v_head_dim) + elif forward_batch.forward_mode.is_decode(): + return self.forward_decode( + q, + k, + v, + layer, + forward_batch, + save_kv_cache=save_kv_cache, + **kwargs, + ) + else: + return self.forward_extend( + q, + k, + v, + layer, + forward_batch, + save_kv_cache=save_kv_cache, + **kwargs, + ) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + ): + """Run a forward for decode.""" + raise NotImplementedError() + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + ): + """Run a forward for extend.""" + raise NotImplementedError() + + def support_triton(self): + """Check if the current backend supports triton.""" + return True diff --git a/python/sglang/srt/layers/attention/cutlass_mla_backend.py b/python/sglang/srt/layers/attention/cutlass_mla_backend.py new file mode 100644 index 000000000..eb0cae262 --- /dev/null +++ b/python/sglang/srt/layers/attention/cutlass_mla_backend.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +""" +Support attention backend for Cutlass MLA. + +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Union + +import torch +import triton + +from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend +from sglang.srt.layers.attention.utils import create_flashmla_kv_indices_triton +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils import is_cuda + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpecInfo + +_is_cuda = is_cuda() +if _is_cuda: + from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size + + +# Cutlass MLA only supports pagesize=128 +PAGE_SIZE = 128 + + +@dataclass +class CutlassMLADecodeMetadata: + workspace: Optional[torch.Tensor] = None + block_kv_indices: Optional[torch.Tensor] = None + + def __init__( + self, + workspace: Optional[torch.Tensor] = None, + block_kv_indices: Optional[torch.Tensor] = None, + ): + self.workspace = workspace + self.block_kv_indices = block_kv_indices + + +class CutlassMLABackend(FlashInferMLAAttnBackend): + """Cutlass attention kernels.""" + + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + kv_indptr_buf: Optional[torch.Tensor] = None, + kv_last_page_len_buf: Optional[torch.Tensor] = None, + ): + super().__init__( + model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf + ) + + self.num_q_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.num_kv_heads = model_runner.model_config.get_num_kv_heads( + get_attention_tp_size() + ) + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.num_local_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.forward_metadata: Union[CutlassMLADecodeMetadata] = None + self.kv_lora_rank = model_runner.model_config.kv_lora_rank + self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim + self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim + self.v_head_dim = model_runner.model_config.v_head_dim + self.scaling = model_runner.model_config.scaling + self.data_type = model_runner.kv_cache_dtype + self.q_data_type = model_runner.dtype + self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim + + def init_forward_metadata(self, forward_batch: ForwardBatch): + + bs = forward_batch.batch_size + spec_info = forward_batch.spec_info + if forward_batch.forward_mode.is_decode_or_idle(): + if spec_info is None: + max_seqlen_pad = triton.cdiv( + forward_batch.seq_lens_cpu.max().item(), PAGE_SIZE + ) + block_kv_indices = torch.full( + (bs, max_seqlen_pad), + -1, + dtype=torch.int32, + device=forward_batch.seq_lens.device, + ) + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + None, + block_kv_indices, + self.req_to_token.stride(0), + max_seqlen_pad, + PAGED_SIZE=PAGE_SIZE, + ) + workspace_size = cutlass_mla_get_workspace_size( + max_seqlen_pad * PAGE_SIZE, bs, num_kv_splits=1 + ) + workspace = torch.empty( + workspace_size, device="cuda", dtype=torch.uint8 + ) + self.forward_metadata = CutlassMLADecodeMetadata( + workspace, + block_kv_indices, + ) + else: + super().init_forward_metadata(forward_batch) + else: + super().init_forward_metadata(forward_batch) + + def init_cuda_graph_state( + self, + max_bs: int, + max_num_tokens: int, + block_kv_indices: Optional[torch.Tensor] = None, + ): + if block_kv_indices is None: + cuda_graph_kv_indices = torch.full( + (max_bs, (self.max_context_len + PAGE_SIZE) // PAGE_SIZE), + 1, + dtype=torch.int32, + device="cuda", + ) + else: + cuda_graph_kv_indices = block_kv_indices + + workspace_size = cutlass_mla_get_workspace_size( + cuda_graph_kv_indices.shape[1] * PAGE_SIZE, max_bs, num_kv_splits=1 + ) + self.cuda_graph_mla_workspace = torch.empty( + workspace_size, device="cuda", dtype=torch.uint8 + ) + self.cuda_graph_kv_indices = cuda_graph_kv_indices + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + ): + if forward_mode.is_decode_or_idle(): + if spec_info is None: + max_seqlen_pad = self.cuda_graph_kv_indices.shape[1] + + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + None, + self.cuda_graph_kv_indices, + self.req_to_token.stride(0), + self.cuda_graph_kv_indices.stride(0), + PAGED_SIZE=PAGE_SIZE, + ) + self.forward_metadata = CutlassMLADecodeMetadata( + self.cuda_graph_mla_workspace, + self.cuda_graph_kv_indices[:bs, :max_seqlen_pad], + ) + else: + super().init_forward_metadata_capture_cuda_graph( + bs, + num_tokens, + req_pool_indices, + seq_lens, + encoder_lens, + forward_mode, + spec_info, + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + seq_lens_cpu: Optional[torch.Tensor], + ): + + if forward_mode.is_decode_or_idle(): + assert seq_lens_cpu is not None + seq_lens = seq_lens[:bs] + + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices[:bs], + seq_lens, + None, + self.cuda_graph_kv_indices, + self.req_to_token.stride(0), + self.cuda_graph_kv_indices.stride(0), + PAGED_SIZE=PAGE_SIZE, + ) + else: + super().init_forward_metadata_replay_cuda_graph( + bs, + req_pool_indices, + seq_lens, + seq_lens_sum, + encoder_lens, + forward_mode, + spec_info, + seq_lens_cpu, + ) + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + ): + cache_loc = forward_batch.out_cache_loc + + if k is not None: + assert v is not None + if save_kv_cache: + if k_rope is not None: + forward_batch.token_to_kv_pool.set_mla_kv_buffer( + layer, + cache_loc, + k, + k_rope, + ) + else: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, + cache_loc, + k, + v, + ) + + # Reshape inputs + if q_rope is not None: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + else: + reshaped_q = q.view(-1, layer.tp_q_head_num, layer.head_dim) + q_nope = reshaped_q[:, :, : layer.v_head_dim] + q_rope = reshaped_q[:, :, layer.v_head_dim :] + + q_nope = q_nope.to(self.q_data_type) + q_rope = q_rope.to(self.q_data_type) + + k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + + o = cutlass_mla_decode( + q_nope=q_nope, + q_pe=q_rope, + kv_c_and_k_pe_cache=k_cache.view(-1, PAGE_SIZE, self.kv_cache_dim), + seq_lens=forward_batch.seq_lens.to(torch.int32), + page_table=self.forward_metadata.block_kv_indices, + workspace=self.forward_metadata.workspace, + sm_scale=layer.scaling, + num_kv_splits=1, + ) + + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) diff --git a/python/sglang/srt/layers/attention/double_sparsity_backend.py b/python/sglang/srt/layers/attention/double_sparsity_backend.py new file mode 100644 index 000000000..47b867f61 --- /dev/null +++ b/python/sglang/srt/layers/attention/double_sparsity_backend.py @@ -0,0 +1,257 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + + +class DoubleSparseAttnBackend(AttentionBackend): + def __init__(self, model_runner: ModelRunner): + # Lazy import to avoid the initialization of cuda context + from sglang.srt.layers.attention.triton_ops.double_sparsity_attention import ( + extend_attention_fwd, + flash_decode_attention_fwd, + flash_decode_sparse_attention_fwd, + ) + + super().__init__() + + self.decode_attention_fwd = flash_decode_attention_fwd + self.decode_sparse_attention_fwd = flash_decode_sparse_attention_fwd + self.extend_attention_fwd = extend_attention_fwd + self.num_head = model_runner.model_config.num_attention_heads + self.head_dim = model_runner.model_config.hidden_size // self.num_head + self.heavy_token_num = model_runner.server_args.ds_heavy_token_num + + self.sorted_channels = model_runner.sorted_channels + self.sparse_decode_thresold = ( + model_runner.server_args.ds_sparse_decode_threshold + ) + self.att_out_approx: torch.Tensor = None + self.mid_out: torch.Tensor = None + self.mid_o_logexpsum: torch.Tensor = None + + # TODO: Change the hard-coded block_seq_num + self.BLOCK_SEQ = 128 + + if global_server_args_dict.get("triton_attention_reduce_in_fp32", False): + self.reduce_dtype = torch.float32 + else: + self.reduce_dtype = torch.float16 + + self.forward_metadata = None + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init auxiliary variables for triton attention backend.""" + + if forward_batch.forward_mode.is_decode(): + start_loc = torch.zeros_like(forward_batch.seq_lens, dtype=torch.int32) + start_loc[1:] = torch.cumsum(forward_batch.seq_lens[:-1], dim=0) + + total_num_tokens = torch.sum(forward_batch.seq_lens).item() + attn_logits = torch.empty( + (self.num_head, total_num_tokens), + dtype=self.reduce_dtype, + device="cuda", + ) + + max_seq_len = torch.max(forward_batch.seq_lens).item() + min_seq_len = torch.min(forward_batch.seq_lens).item() + max_extend_len = None + # NOTE: Align sequence order with req_to_token order + ds_req_to_token = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices + ] + + bsz = forward_batch.seq_lens.shape[0] + + att_out_approx = torch.empty( + [self.num_head, bsz, max_seq_len], + dtype=self.reduce_dtype, + device="cuda", + ) + + block_seq_num = ( + self.heavy_token_num + self.BLOCK_SEQ - 1 + ) // self.BLOCK_SEQ + + mid_out = torch.empty( + [bsz, self.num_head, block_seq_num, self.head_dim], + dtype=torch.float32, + device="cuda", + ) + mid_o_logexpsum = torch.empty( + [bsz, self.num_head, block_seq_num], dtype=torch.float32, device="cuda" + ) + self.att_out_approx = att_out_approx + self.mid_out = mid_out + self.mid_o_logexpsum = mid_o_logexpsum + + else: + start_loc = attn_logits = max_seq_len = min_seq_len = None + prefix_lens = forward_batch.extend_prefix_lens + max_extend_len = torch.max(forward_batch.seq_lens - prefix_lens).item() + ds_req_to_token = None + + self.forward_metadata = ( + start_loc, + attn_logits, + max_seq_len, + min_seq_len, + max_extend_len, + ds_req_to_token, + ) + + def forward_extend( + self, + q, + k, + v, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + # TODO: reuse the buffer across layers + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + k_label = torch.gather( + k, + 2, + self.sorted_channels[layer.layer_id] + .unsqueeze(0) + .expand(k.shape[0], -1, -1), + ) + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v, k_label + ) + + ( + start_loc, + attn_logits, + max_seq_len, + min_seq_len, + max_extend_len, + ds_req_to_token, + ) = self.forward_metadata + self.extend_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + k.contiguous(), + v.contiguous(), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + forward_batch.req_to_token_pool.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.extend_seq_lens, + forward_batch.extend_start_loc, + max_extend_len, + layer.scaling, + layer.logit_cap, + ) + return o + + def forward_decode( + self, + q, + k, + v, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + # During torch.compile, there is a bug in rotary_emb that causes the + # output value to have a 3D tensor shape. This reshapes the output correctly. + q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim) + + # TODO: reuse the buffer across layers + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + # TODO: Add min seqlen + ( + start_loc, + attn_logits, + max_seq_len, + min_seq_len, + max_extend_len, + ds_req_to_token, + ) = self.forward_metadata + + k_label = torch.gather( + k, + 2, + self.sorted_channels[layer.layer_id] + .unsqueeze(0) + .expand(k.shape[0], -1, -1), + ) + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v, k_label + ) + + # NOTE(Andy) shouldn't be used when max_len_in_batch < heavy_token_num + # and set a minimum value for sparse_decode + if ( + min_seq_len < self.heavy_token_num + or max_seq_len < self.sparse_decode_thresold + ): + self.decode_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + forward_batch.req_to_token_pool.req_to_token, + forward_batch.req_pool_indices, + start_loc, + forward_batch.seq_lens, + attn_logits, + max_seq_len, + layer.scaling, + layer.logit_cap, + ) + else: + # TODO(Andy): indexing with torch.gather or torch.index_select or customized kernel + q_label = torch.gather( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + 2, + self.sorted_channels[layer.layer_id] + .unsqueeze(0) + .expand(q.shape[0], -1, -1), + ) + self.decode_sparse_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + o.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + q_label, + forward_batch.token_to_kv_pool.get_label_buffer(layer.layer_id), + ds_req_to_token, + forward_batch.seq_lens, + max_seq_len, + layer.scaling, + layer.logit_cap, + self.heavy_token_num, + self.att_out_approx, + self.mid_out, + self.mid_o_logexpsum, + self.BLOCK_SEQ, + ) + + return o diff --git a/python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py b/python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py new file mode 100644 index 000000000..84876b438 --- /dev/null +++ b/python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py @@ -0,0 +1,1700 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Attention layer with Dual chunk flash attention and sparse attention. +""" +import functools +import logging +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import torch +import torch.nn.functional as F +from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache +from sgl_kernel.sparse_flash_attn import ( + convert_vertical_slash_indexes, + convert_vertical_slash_indexes_mergehead, + sparse_attn_func, +) + +from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_rank +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.flashattention_backend import FlashAttentionMetadata +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + + +logger = logging.getLogger(__name__) + + +@dataclass +class DualChunkFlashAttentionMetadata: + """Metadata for FlashAttentionBackend. + + NOTE: Any python object stored here is not updated when it is + cuda-graph replayed. If you have values that need to be changed + dynamically, it should be stored in tensor. The tensor has to be + updated from `CUDAGraphRunner.forward` API. + """ + + # (batch_size,). The sequence length per sequence. Sequence length means + # the computed tokens + new tokens None if it is a decoding. + seq_lens: Optional[List[int]] = None + # seq_lens stored as a tensor. + seq_lens_tensor: Optional[torch.Tensor] = None + # Maximum sequence length among prefill batch. 0 if there are decoding + # requests only. + max_seq_len: int = None + + # (batch_size,). The orig sequence length per sequence. + orig_seq_lens: Optional[List[int]] = None + + # orig_seq_lens stored as a tensor. + orig_seq_lens_tensor: Optional[torch.Tensor] = None + + # Block addresses per sequence. (Seq id -> list of physical block) + # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks + # in the kv cache. Each block can contain up to block_size tokens. + # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph + # captured. + block_tables: Optional[torch.Tensor] = None + + # (batch_size + 1,). The cumulative subquery lengths of the sequences in + # the batch, used to index into subquery. E.g., if the subquery length + # is [4, 6], it is [0, 4, 10]. + query_start_loc: Optional[torch.Tensor] = None + # (batch_size + 1,). The cumulative sequence lengths of the sequences in + # the batch, used to index into sequence. E.g., if the sequence length is + # [4, 6], it is [0, 4, 10]. + seq_start_loc: Optional[torch.Tensor] = None + + # Length scaling factor + scaling_factor: Optional[torch.Tensor] = None + + # (batch_size,). Sequence lengths for intra attention. + seq_lens_intra: Optional[torch.Tensor] = None + + # Max sequence length for intra attention. + max_seq_len_intra: Optional[int] = None + + # (batch_size, num_blocks). Block table for intra attention. + block_tables_intra: Optional[torch.Tensor] = None + + # (batch_size,). Sequence lengths for succ attention. + seq_lens_succ: Optional[torch.Tensor] = None + + # Max sequence length for succ attention. + max_seq_len_succ: Optional[int] = None + + # (batch_size, num_blocks). Block table for succ attention. + block_tables_succ: Optional[torch.Tensor] = None + + # (batch_size,). Sequence lengths for inter attention. + seq_lens_inter: Optional[torch.Tensor] = None + + # Max sequence length for inter attention. + max_seq_len_inter: Optional[int] = None + + +class DualChunkFlashAttentionBackend(AttentionBackend): + def __init__( + self, + model_runner: "ModelRunner", + ) -> None: + self.forward_metadata: FlashAttentionMetadata = None + self.device = model_runner.device + self.max_context_len = model_runner.model_config.context_len + self.num_heads = model_runner.model_config.get_num_attention_heads( + model_runner.server_args.tp_size + ) + self.num_kv_heads = model_runner.model_config.get_num_kv_heads( + model_runner.server_args.tp_size + ) + self.head_size = model_runner.model_config.head_dim + + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.kv_cache_dtype = model_runner.kv_cache_dtype + self.kv_cache_dtype_str = model_runner.server_args.kv_cache_dtype + self.page_size = model_runner.page_size + + assert self.num_heads % self.num_kv_heads == 0 + self.num_queries_per_kv = self.num_heads // self.num_kv_heads + + dual_chunk_attention_config = getattr( + model_runner.model_config.hf_config, "dual_chunk_attention_config", None + ) + assert dual_chunk_attention_config is not None + self.chunk_size = dual_chunk_attention_config.get("chunk_size", 8192) + self.local_size = dual_chunk_attention_config.get("local_size", 1024) + self.original_max_position_embeddings = dual_chunk_attention_config.get( + "original_max_position_embeddings", 0 + ) + self.sparse_attention_config = dual_chunk_attention_config.get( + "sparse_attention_config", None + ) + if not self.sparse_attention_config: + logger.warning_once( + "Sparse attention will not be enabled as " + "sparse attention config is not provided." + ) + self.sparse_attention_enabled = dual_chunk_attention_config.get( + "sparse_attention_enabled", self.sparse_attention_config is not None + ) + self.sparse_attention_threshold = dual_chunk_attention_config.get( + "sparse_attention_threshold", 32768 + ) + self.sparse_attention_last_q = dual_chunk_attention_config.get( + "sparse_attention_last_q", 64 + ) + self.dual_chunk_attention_config = dual_chunk_attention_config + + if self.sparse_attention_enabled: + self.arange = torch.arange(self.sparse_attention_last_q, device="cuda") + self.last_q_mask = ( + self.arange[None, None, :, None] >= self.arange[None, None, None, :] + ) + + @functools.lru_cache() + def get_sparse_attention_config(self, layer_idx) -> List[Dict[str, Any]]: + layer_sparse_attention_config = { + int(i): j for i, j in self.sparse_attention_config[layer_idx].items() + } + start_head = self.num_heads * get_tensor_model_parallel_rank() + end_head = start_head + self.num_heads + return [layer_sparse_attention_config[i] for i in range(start_head, end_head)] + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Initialize forward metadata hence all layers in the forward pass can reuse it.""" + + forward_mode: ForwardMode = forward_batch.forward_mode + assert forward_mode.is_prefill() or forward_mode.is_decode() + batch_size = forward_batch.batch_size + + metadata = DualChunkFlashAttentionMetadata() + metadata.seq_lens_tensor = forward_batch.seq_lens.to(torch.int32) + metadata.seq_lens = forward_batch.seq_lens.tolist() + metadata.max_seq_len = forward_batch.seq_lens.max().item() + + metadata.orig_seq_lens_tensor = forward_batch.orig_seq_lens + metadata.orig_seq_lens = forward_batch.orig_seq_lens.tolist() + + metadata.block_tables = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len + ] + # Convert the block table to a strided format. + if self.page_size > 1: + strided_indices = torch.arange( + 0, metadata.block_tables.shape[1], self.page_size, device=self.device + ) + metadata.block_tables = ( + metadata.block_tables[:, strided_indices] // self.page_size + ) + + metadata.query_start_loc = torch.zeros( + batch_size + 1, dtype=torch.int32, device=metadata.seq_lens_tensor.device + ) + if forward_mode.is_prefill(): + metadata.query_start_loc[1:] = torch.cumsum( + forward_batch.extend_seq_lens.to(torch.int32), dim=0, dtype=torch.int32 + ) + else: + metadata.query_start_loc[1:] = torch.cumsum( + torch.arange( + batch_size, + dtype=metadata.query_start_loc.dtype, + device=metadata.query_start_loc.device, + ), + dim=0, + dtype=torch.int32, + ) + metadata.seq_start_loc = torch.zeros( + batch_size + 1, dtype=torch.int32, device=metadata.seq_lens_tensor.device + ) + metadata.seq_start_loc[1:] = torch.cumsum( + metadata.seq_lens_tensor, dim=0, dtype=torch.int32 + ) + + if self.original_max_position_embeddings > 0: + if forward_mode.is_prefill(): + metadata.scaling_factor = ( + 0.1 + * torch.log( + metadata.orig_seq_lens_tensor + / self.original_max_position_embeddings + ) + + 1.0 + ).clip(min=1) + else: + metadata.scaling_factor = ( + 0.1 + * torch.log( + metadata.orig_seq_lens_tensor + / self.original_max_position_embeddings + ) + + 1.0 + ).clip(min=1) + + if forward_mode.is_decode(): + cache_seq_lens = metadata.orig_seq_lens_tensor + + chunk_len = self.chunk_size - self.local_size + chunk_num_curr = (cache_seq_lens - 1) // chunk_len + + seq_lens_intra = cache_seq_lens - chunk_num_curr * chunk_len + max_seq_len_intra = seq_lens_intra.max().item() + metadata.seq_lens_intra = seq_lens_intra + metadata.max_seq_len_intra = max_seq_len_intra + + block_tables_intra = torch.zeros( + batch_size, + (max_seq_len_intra - 1) // self.page_size + 1, + dtype=metadata.block_tables.dtype, + device=metadata.block_tables.device, + ) + for i in range(batch_size): + st = chunk_num_curr[i] * chunk_len // self.page_size + ed = min( + st + (max_seq_len_intra - 1) // self.page_size + 1, + (cache_seq_lens[i] - 1) // self.page_size + 1, + ) + block_tables_intra[i, : ed - st] = metadata.block_tables[i, st:ed] + metadata.block_tables_intra = block_tables_intra + + metadata.seq_lens_succ = ( + chunk_num_curr - (chunk_num_curr - 1).clip(min=0) + ) * chunk_len + metadata.max_seq_len_succ = metadata.seq_lens_succ.max().item() + if metadata.max_seq_len_succ: + block_tables_succ = torch.zeros( + batch_size, + (metadata.max_seq_len_succ - 1) // self.page_size + 1, + dtype=metadata.block_tables.dtype, + device=metadata.block_tables.device, + ) + for i in range(batch_size): + start = ( + (chunk_num_curr[i] - 1).clip(min=0) + * chunk_len + // self.page_size + ) + end = min( + start + (metadata.max_seq_len_succ - 1) // self.page_size + 1, + (cache_seq_lens[i] - 1) // self.page_size + 1, + ) + block_tables_succ[i, : end - start] = metadata.block_tables[ + i, start:end + ] + metadata.block_tables_succ = block_tables_succ + + metadata.seq_lens_inter = (chunk_num_curr - 1).clip(min=0) * chunk_len + metadata.max_seq_len_inter = metadata.seq_lens_inter.max().item() + + self.forward_metadata = metadata + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: "RadixAttention", + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + # Use precomputed metadata across all layers + metadata = self.forward_metadata + + ( + query, + query_succ, + query_inter, + query_succ_critical, + query_inter_critical, + ) = torch.split(q, q.shape[-1] // 5, dim=-1) + + # Reshape the query, key, and value tensors. + query = query.view(-1, self.num_heads, self.head_size) + query_succ = query_succ.view(-1, self.num_heads, self.head_size) + query_inter = query_inter.view(-1, self.num_heads, self.head_size) + query_succ_critical = query_succ_critical.view( + -1, self.num_heads, self.head_size + ) + query_inter_critical = query_inter_critical.view( + -1, self.num_heads, self.head_size + ) + key = k.view(-1, self.num_kv_heads, self.head_size) + value = v.view(-1, self.num_kv_heads, self.head_size) + + # apply DCA scaling + if self.original_max_position_embeddings > 0: + assert metadata.scaling_factor is not None + assert metadata.query_start_loc is not None + assert metadata.orig_seq_lens is not None + current_start = 0 + query_start_loc_cpu = metadata.query_start_loc.cpu() + for i in range(len(metadata.orig_seq_lens)): + current_end = ( + current_start + + (query_start_loc_cpu[i + 1] - query_start_loc_cpu[i]).item() + ) + key[current_start:current_end].mul_(metadata.scaling_factor[i]) + current_start = current_end + assert current_end <= self.max_context_len + + # Do multi-head attention + key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer( + layer.layer_id + ) + key_cache = key_cache.view( + -1, self.page_size, layer.tp_k_head_num, layer.head_dim + ) + value_cache = value_cache.view( + -1, self.page_size, layer.tp_v_head_num, layer.head_dim + ) + + if key is not None and value is not None: + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, + forward_batch.out_cache_loc, + key, + value, + layer.k_scale, + layer.v_scale, + ) + + if not save_kv_cache: + # profile run + o = flash_attn_varlen_func( + q=query, + k=key, + v=value, + cu_seqlens_q=metadata.seq_start_loc, + cu_seqlens_k=metadata.seq_start_loc, + max_seqlen_q=metadata.max_seq_len, + max_seqlen_k=metadata.max_seq_len, + softmax_scale=layer.scaling, + causal=True, + ) + else: + # prefill/chunked-prefill + # get per layer sparse attention config + if self.sparse_attention_enabled: + self.layer_sparse_attention_config = self.get_sparse_attention_config( + layer.layer_id + ) + assert metadata.orig_seq_lens is not None + o = self._dual_chunk_flash_attn_prefill( + q=query, + q_succ=query_succ, + q_inter=query_inter, + q_succ_critical=query_succ_critical, + q_inter_critical=query_inter_critical, + k=key_cache, + v=value_cache, + cu_seqlens_q=metadata.query_start_loc, + cu_seqlens_k=metadata.seq_start_loc, + orig_seq_lens=metadata.orig_seq_lens, + scaling_factor=metadata.scaling_factor, + softmax_scale=layer.scaling, + causal=True, + window_size=(-1, -1), + block_table=metadata.block_tables, + chunk_size=self.chunk_size, + local_size=self.local_size, + ) + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: "RadixAttention", + forward_batch: ForwardBatch, + save_kv_cache=True, + ) -> torch.Tensor: + # Use precomputed metadata across all layers + metadata = self.forward_metadata + + ( + query, + query_succ, + query_inter, + query_succ_critical, + query_inter_critical, + ) = torch.split(q, q.shape[-1] // 5, dim=-1) + + # Reshape the query, key, and value tensors. + query = query.view(-1, self.num_heads, self.head_size) + query_succ = query_succ.view(-1, self.num_heads, self.head_size) + query_inter = query_inter.view(-1, self.num_heads, self.head_size) + query_succ_critical = query_succ_critical.view( + -1, self.num_heads, self.head_size + ) + query_inter_critical = query_inter_critical.view( + -1, self.num_heads, self.head_size + ) + key = k.view(-1, self.num_kv_heads, self.head_size) + value = v.view(-1, self.num_kv_heads, self.head_size) + + key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer( + layer.layer_id + ) + key_cache = key_cache.view( + -1, self.page_size, layer.tp_k_head_num, layer.head_dim + ) + value_cache = value_cache.view( + -1, self.page_size, layer.tp_v_head_num, layer.head_dim + ) + + if key is not None and value is not None: + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, + forward_batch.out_cache_loc, + key, + value, + layer.k_scale, + layer.v_scale, + ) + + # apply DCA scaling + if self.original_max_position_embeddings > 0: + assert metadata.scaling_factor is not None + scaling_factor = metadata.scaling_factor + key.mul_(scaling_factor.unsqueeze(-1).unsqueeze(-1)) + + o = self._dual_chunk_flash_attn_decoding( + query.unsqueeze(1), + query_succ.unsqueeze(1), + query_inter.unsqueeze(1), + key_cache, + value_cache, + block_table=metadata.block_tables, + cache_seqlens=metadata.seq_lens_tensor, + softmax_scale=layer.scaling, + causal=True, + chunk_size=self.chunk_size, + local_size=self.local_size, + original_max_position_embeddings=self.original_max_position_embeddings, + decode_meta=metadata, + ).squeeze(1) + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + """Initialize CUDA graph state for the attention backend. + + Args: + max_bs (int): Maximum batch size to support in CUDA graphs + + This creates fixed-size tensors that will be reused during CUDA graph replay + to avoid memory allocations. + """ + self.decode_metadata = { + "seq_lens_tensor": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "orig_seq_lens_tensor": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "scaling_factor": torch.zeros( + max_bs, dtype=torch.float32, device=self.device + ), + "block_tables": torch.zeros( + max_bs, + (self.max_context_len - 1) // self.page_size + 1, + dtype=torch.int32, + device=self.device, + ), + "block_tables_intra": torch.zeros( + max_bs, + (self.max_context_len - 1) // self.page_size + 1, + dtype=torch.int32, + device=self.device, + ), + "seq_lens_intra": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "block_tables_succ": torch.zeros( + max_bs, + (self.max_context_len - 1) // self.page_size + 1, + dtype=torch.int32, + device=self.device, + ), + "seq_lens_succ": torch.zeros(max_bs, dtype=torch.int32, device=self.device), + "seq_lens_inter": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + } + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[None], + ): + metadata = DualChunkFlashAttentionMetadata() + + if forward_mode.is_decode_or_idle(): + if self.original_max_position_embeddings > 0: + metadata.scaling_factor = self.decode_metadata["scaling_factor"][:bs] + + metadata.seq_lens_tensor = self.decode_metadata["seq_lens_tensor"][:bs] + metadata.orig_seq_lens_tensor = self.decode_metadata[ + "orig_seq_lens_tensor" + ][:bs] + metadata.max_seq_len = self.max_context_len + metadata.block_tables = self.decode_metadata["block_tables"][ + req_pool_indices, : + ] + + # intra + metadata.max_seq_len_intra = self.max_context_len + metadata.seq_lens_intra = self.decode_metadata["seq_lens_intra"][:bs] + + metadata.block_tables_intra = self.decode_metadata["block_tables_intra"][ + :bs, : + ] + + # succ + metadata.seq_lens_succ = self.decode_metadata["seq_lens_succ"][:bs] + metadata.max_seq_len_succ = self.max_context_len + + metadata.block_tables_succ = self.decode_metadata["block_tables_succ"][ + :bs, : + ] + + metadata.seq_lens_inter = self.decode_metadata["seq_lens_inter"][:bs] + metadata.max_seq_len_inter = self.max_context_len + + self.decode_metadata[bs] = metadata + + self.forward_metadata = metadata + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[None], + seq_lens_cpu: Optional[torch.Tensor], + out_cache_loc: torch.Tensor = None, + ): + """Initialize forward metadata for replaying CUDA graph.""" + assert forward_mode.is_decode() + seq_lens = seq_lens[:bs] + req_pool_indices = req_pool_indices[:bs] + metadata = self.decode_metadata[bs] + + metadata.seq_lens_tensor.copy_(seq_lens.to(torch.int32)) + metadata.seq_lens = seq_lens.tolist() + metadata.max_seq_len = seq_lens.max().item() + + metadata.orig_seq_lens_tensor.copy_(seq_lens) + metadata.orig_seq_lens = seq_lens.tolist() + + block_tables = self.req_to_token[req_pool_indices, : metadata.max_seq_len] + # Convert the block table to a strided format. + if self.page_size > 1: + strided_indices = torch.arange( + 0, block_tables.shape[1], self.page_size, device=self.device + ) + block_tables = block_tables[:, strided_indices] // self.page_size + metadata.block_tables.fill_(0) + metadata.block_tables[: block_tables.shape[0], : block_tables.shape[1]].copy_( + block_tables + ) + + if self.original_max_position_embeddings > 0: + scaling_factor = ( + 0.1 + * torch.log( + metadata.orig_seq_lens_tensor + / self.original_max_position_embeddings + ) + + 1.0 + ).clip(min=1) + metadata.scaling_factor.copy_(scaling_factor) + + cache_seq_lens = metadata.orig_seq_lens_tensor + + chunk_len = self.chunk_size - self.local_size + chunk_num_curr = (cache_seq_lens - 1) // chunk_len + + seq_lens_intra = cache_seq_lens - chunk_num_curr * chunk_len + max_seq_len_intra = seq_lens_intra.max().item() + metadata.seq_lens_intra.copy_(seq_lens_intra) + metadata.max_seq_len_intra = max_seq_len_intra + + metadata.block_tables_intra.fill_(0) + for i in range(bs): + st = chunk_num_curr[i] * chunk_len // self.page_size + ed = min( + st + (max_seq_len_intra - 1) // self.page_size + 1, + (cache_seq_lens[i] - 1) // self.page_size + 1, + ) + metadata.block_tables_intra[i, : ed - st] = metadata.block_tables[i, st:ed] + + seq_lens_succ = (chunk_num_curr - (chunk_num_curr - 1).clip(min=0)) * chunk_len + metadata.seq_lens_succ.copy_(seq_lens_succ) + metadata.max_seq_len_succ = metadata.seq_lens_succ.max().item() + if metadata.max_seq_len_succ: + metadata.block_tables_succ.fill_(0) + for i in range(bs): + start = ( + (chunk_num_curr[i] - 1).clip(min=0) * chunk_len // self.page_size + ) + end = min( + start + (metadata.max_seq_len_succ - 1) // self.page_size + 1, + (cache_seq_lens[i] - 1) // self.page_size + 1, + ) + metadata.block_tables_succ[i, : end - start] = metadata.block_tables[ + i, start:end + ] + + seq_lens_inter = (chunk_num_curr - 1).clip(min=0) * chunk_len + metadata.seq_lens_inter.copy_(seq_lens_inter) + metadata.max_seq_len_inter = metadata.seq_lens_inter.max().item() + + self.forward_metadata = metadata + + def get_cuda_graph_seq_len_fill_value(self): + """Get the fill value for sequence length in CUDA graph.""" + return 1 + + def _dual_chunk_flash_attn_prefill( + self, + q, + q_succ, + q_inter, + q_succ_critical, + q_inter_critical, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + orig_seq_lens: List[int], + scaling_factor: torch.Tensor, + softmax_scale: float, + causal: Optional[bool] = True, + window_size: Tuple[int, int] = (-1, -1), + block_table: Optional[torch.Tensor] = None, + chunk_size: int = 8192, + local_size: int = 1024, + ): + if not causal: + raise ValueError("Dual Chunk Attention does not support causal=False") + if window_size != (-1, -1): + raise ValueError("Dual Chunk Attention does not support window_size") + + cu_seqlens_q_cpu = cu_seqlens_q.cpu().tolist() + cu_seqlens_k_cpu = cu_seqlens_k.cpu().tolist() + all_outputs = [] + + for i in range(0, len(cu_seqlens_q_cpu) - 1): + qs = cu_seqlens_q_cpu[i] + qe = cu_seqlens_q_cpu[i : i + 2][-1] + ks = cu_seqlens_k_cpu[i] + ke = cu_seqlens_k_cpu[i : i + 2][-1] + + current_q = q[qs:qe] + current_q_succ = q_succ[qs:qe] + current_q_inter = q_inter[qs:qe] + current_q_succ_critical = q_succ_critical[qs:qe] + current_q_inter_critical = q_inter_critical[qs:qe] + + if block_table is None: + current_k = k[ks:ke] + current_v = v[ks:ke] + current_block_table = None + current_orig_seq_len = orig_seq_lens[i] + else: + current_block_table = block_table[i] + current_orig_seq_len = orig_seq_lens[i] + current_k = k + current_v = v + sparse_attn_enabled = ( + self.sparse_attention_enabled + and current_orig_seq_len > self.sparse_attention_threshold + ) + + if current_q.shape[0] == 0: + continue + + if current_k.shape[0] == 0: + all_outputs.append( + torch.zeros( + (current_q.shape[0], current_q.shape[1], v.shape[2]), + device=q.device, + dtype=q.dtype, + ) + ) + continue + + current_output = torch.empty_like(current_q) + group_size = int(current_q.size(-2) / current_k.size(-2)) + + if sparse_attn_enabled: + num_device_q_heads = current_q.size(-2) + heads_vertical_size = torch.empty( + size=(num_device_q_heads,), dtype=torch.int32 + ) + heads_slash_size = torch.empty( + size=(num_device_q_heads,), dtype=torch.int32 + ) + for head_id in range(current_q.size(-2)): + ( + ty, + vertical_size, + slash_size, + _, + ) = self.layer_sparse_attention_config[head_id] + assert ty == "vertical_and_slash", "only support slash mode" + + if vertical_size == 30: + vertical_size += 100 + heads_vertical_size[head_id] = vertical_size + heads_slash_size[head_id] = slash_size + + current_output = self._dual_chunk_flash_attn_prefill_func( + current_q, # allheads + current_q_succ, + current_q_inter, + current_q_succ_critical, + current_q_inter_critical, + current_k, + current_v, + current_block_table, + softmax_scale, + chunk_size, + local_size, + scaling_factor[i].item(), + ke - ks, + sparse_attn_enabled=sparse_attn_enabled, + heads_vertical_size=heads_vertical_size, + heads_slash_size=heads_slash_size, + group_size=group_size, + ) + else: + for head_id in range(current_q.size(-2)): + # (seq_len, num_heads, head_size) + current_q_head = current_q[:, head_id, :].unsqueeze(1) + current_q_succ_head = current_q_succ[:, head_id, :].unsqueeze(1) + current_q_inter_head = current_q_inter[:, head_id, :].unsqueeze(1) + current_q_succ_head_critical = current_q_succ_critical[ + :, head_id, : + ].unsqueeze(1) + current_q_inter_head_critical = current_q_inter_critical[ + :, head_id, : + ].unsqueeze(1) + if block_table is not None: + current_k_head = current_k[ + ..., head_id // group_size, : + ].unsqueeze(2) + current_v_head = current_v[ + ..., head_id // group_size, : + ].unsqueeze(2) + + else: + current_k_head = current_k[:, head_id, :].unsqueeze(1) + current_v_head = current_v[:, head_id, :].unsqueeze(1) + + current_out = self._dual_chunk_flash_attn_prefill_func( + current_q_head, + current_q_succ_head, + current_q_inter_head, + current_q_succ_head_critical, + current_q_inter_head_critical, + current_k_head, + current_v_head, + current_block_table, + softmax_scale, + chunk_size, + local_size, + scaling_factor[i].item(), + ke - ks, + sparse_attn_enabled=sparse_attn_enabled, + ) + current_output[:, head_id : head_id + 1, :] = current_out + all_outputs.append(current_output) + return torch.cat(all_outputs, dim=0) + + def _dual_chunk_flash_attn_prefill_func( + self, + q, + q_succ, + q_inter, + q_succ_critical, + q_inter_critical, + k, + v, + block_table, + softmax_scale: float, + chunk_size: int, + local_size: int, + scaling_factor: float, + k_length: int, + sparse_attn_enabled: Optional[bool] = True, + heads_vertical_size=None, + heads_slash_size=None, + group_size=None, + ): + flash_results = [] + chunk_len = chunk_size - local_size + + if block_table is not None: + block_size = v.shape[1] + if chunk_len % block_size != 0: + raise ValueError("chunk_len must be divisible by block_size.") + else: + block_size = 1 + + if self.original_max_position_embeddings > 0: + softmax_scale = softmax_scale * scaling_factor + + begin = k_length - q.shape[0] + while begin < k_length: + flash_per_chunk = [] + + prev_chunk_end_pos = (begin // chunk_len) * chunk_len + next_chunk_end_pos = prev_chunk_end_pos + chunk_len + end = min(next_chunk_end_pos, k_length) + qbegin = begin - (k_length - q.shape[0]) + qend = end - (k_length - q.shape[0]) + + qk_chunks = [] + q_states_intra = q[qbegin:qend] + # choose critical token + if block_table is not None: + block_tables_intra = _get_block( + block_table, block_size, prev_chunk_end_pos, end + ) + k_states_intra = k[block_tables_intra].view(-1, *k.shape[-2:])[ + : (end - prev_chunk_end_pos) + ] + v_states_intra = v[block_tables_intra].view(-1, *v.shape[-2:])[ + : (end - prev_chunk_end_pos) + ] + else: + block_tables_intra = None + k_states_intra = k[prev_chunk_end_pos:end] + v_states_intra = v[prev_chunk_end_pos:end] + + if sparse_attn_enabled: + last_q_size = min(qend - qbegin, self.sparse_attention_last_q) + _, num_device_k_heads, head_dim = k_states_intra.shape + k_states_intra = ( + k_states_intra.unsqueeze(2) + .repeat(1, 1, group_size, 1) + .reshape(-1, num_device_k_heads * group_size, head_dim) + ) + v_states_intra = ( + v_states_intra.unsqueeze(2) + .repeat(1, 1, group_size, 1) + .reshape(-1, num_device_k_heads * group_size, head_dim) + ) + qk_chunks.append( + (q_states_intra.transpose(0, 1)[:, -last_q_size:] * softmax_scale) + @ k_states_intra.permute(1, 2, 0) + ) + + if prev_chunk_end_pos - chunk_len >= 0: + q_states_succ = q_succ[qbegin:qend] + q_states_succ_critical = q_succ_critical[qbegin:qend] + if block_table is not None: + block_tables_succ = _get_block( + block_table, + block_size, + prev_chunk_end_pos - chunk_len, + prev_chunk_end_pos, + ) + k_states_succ = k[block_tables_succ].view(-1, *k.shape[-2:])[ + :chunk_len + ] + v_states_succ = v[block_tables_succ].view(-1, *v.shape[-2:])[ + :chunk_len + ] + else: + k_states_succ = k[ + prev_chunk_end_pos - chunk_len : prev_chunk_end_pos + ] + v_states_succ = v[ + prev_chunk_end_pos - chunk_len : prev_chunk_end_pos + ] + + if sparse_attn_enabled: + k_states_succ = ( + k_states_succ.unsqueeze(2) + .repeat(1, 1, group_size, 1) + .reshape(-1, num_device_k_heads * group_size, head_dim) + ) + v_states_succ = ( + v_states_succ.unsqueeze(2) + .repeat(1, 1, group_size, 1) + .reshape(-1, num_device_k_heads * group_size, head_dim) + ) + qk_chunks.append( + ( + q_states_succ_critical.transpose(0, 1)[:, -last_q_size:] + * softmax_scale + ) + @ k_states_succ.permute(1, 2, 0) + ) + + if prev_chunk_end_pos - chunk_len * 2 >= 0: + q_states_inter = q_inter[qbegin:qend] + q_states_inter_critical = q_inter_critical[qbegin:qend] + if block_table is not None: + block_tables_inter = _get_block( + block_table, block_size, 0, prev_chunk_end_pos - chunk_len + ) + k_states_inter = k[block_tables_inter].view(-1, *k.shape[-2:])[ + : (prev_chunk_end_pos - chunk_len) + ] + v_states_inter = v[block_tables_inter].view(-1, *v.shape[-2:])[ + : (prev_chunk_end_pos - chunk_len) + ] + else: + k_states_inter = k[: prev_chunk_end_pos - chunk_len] + v_states_inter = v[: prev_chunk_end_pos - chunk_len] + + if sparse_attn_enabled: + k_states_inter = ( + k_states_inter.unsqueeze(2) + .repeat(1, 1, group_size, 1) + .reshape(-1, num_device_k_heads * group_size, head_dim) + ) + v_states_inter = ( + v_states_inter.unsqueeze(2) + .repeat(1, 1, group_size, 1) + .reshape(-1, num_device_k_heads * group_size, head_dim) + ) + qk_chunks.append( + ( + q_states_inter_critical.transpose(0, 1)[:, -last_q_size:] + * softmax_scale + ) + @ k_states_inter.permute(1, 2, 0) + ) + + if sparse_attn_enabled: + reversed_qk = qk_chunks[::-1] + qk = torch.cat(reversed_qk, dim=-1) + + qk[:, :, -last_q_size:] = torch.where( + self.last_q_mask[..., -last_q_size:, -last_q_size:].to(qk.device), + qk[:, :, -last_q_size:], + -torch.inf, + ) + qk = F.softmax(qk, dim=-1, dtype=torch.float32) + + vertical = qk.sum(-2, keepdim=True) + vertical[..., :30] = torch.inf + + # Avoid sorting by using the min/max ints to fill the indexer + # buffers. + int32_max = torch.iinfo(torch.int32).max + int32_min = torch.iinfo(torch.int32).min + n_heads = qk.size()[0] + max_slash_topk = torch.max(heads_slash_size).item() + max_vertical_topk = torch.max(heads_vertical_size).item() + # store each head's slash topk, vertical topk + vertical = vertical.reshape((n_heads, -1)) + # prevent out of range when prompt size < max_vertical_topk + max_vertical_topk = min(vertical.shape[-1], max_vertical_topk) + vertical_topk_buffer = torch.topk( + vertical, max_vertical_topk, -1 + ).indices + slash_topk_buffer = torch.empty( + size=(n_heads, max_slash_topk), dtype=torch.int64, device=qk.device + ) + for head_i in range(n_heads): + # (nqheads=1, lastq, k_len) + head_score = qk[head_i : head_i + 1, :, :] + slash_scores = _sum_all_diagonal_matrix(head_score) + if head_score.size(1) != 1: + # drop right up corner + slash_scores = slash_scores[..., : -last_q_size + 1] + slash_scores[..., -100:] = torch.inf + + head_slash_size = heads_slash_size[head_i] + head_slash_size = min(head_slash_size, vertical.size(-1)) + slash_topk = torch.topk(slash_scores, head_slash_size, -1).indices + # （nheads, max_topk） + slash_topk_buffer[head_i, :head_slash_size] = slash_topk + + # reset heads topk + heads_slash_size[head_i] = head_slash_size + heads_vertical_size[head_i] = min( + heads_vertical_size[head_i], max_vertical_topk + ) + + # store + vertical_buffer = torch.full( + (n_heads, max_vertical_topk), + int32_max, + dtype=torch.int64, + device=q.device, + ) + slash_buffer = torch.full( + (n_heads, max_slash_topk), + int32_min, + dtype=torch.int64, + device=q.device, + ) + succ_vertical_buffer = torch.full( + (n_heads, max_vertical_topk), + int32_max, + dtype=torch.int64, + device=q.device, + ) + succ_slash_buffer = torch.full( + (n_heads, max_slash_topk), + int32_min, + dtype=torch.int64, + device=q.device, + ) + inter_vertical_buffer = torch.full( + (n_heads, max_vertical_topk), + int32_max, + dtype=torch.int64, + device=q.device, + ) + inter_slash_buffer = torch.full( + (n_heads, max_slash_topk), + int32_min, + dtype=torch.int64, + device=q.device, + ) + + vertical_size_buffer = torch.empty( + size=(n_heads,), dtype=torch.int32, device=q.device + ) + slash_sizes_buffer = torch.empty( + size=(n_heads,), dtype=torch.int32, device=q.device + ) + succ_vertical_size_buffer = torch.empty( + size=(n_heads,), dtype=torch.int32, device=q.device + ) + succ_slash_sizes_buffer = torch.empty( + size=(n_heads,), dtype=torch.int32, device=q.device + ) + inter_vertical_size_buffer = torch.empty( + size=(n_heads,), dtype=torch.int32, device=q.device + ) + inter_slash_sizes_buffer = torch.empty( + size=(n_heads,), dtype=torch.int32, device=q.device + ) + + for head_i in range(n_heads): + vertical_topk = vertical_topk_buffer[ + head_i, : heads_vertical_size[head_i] + ] + # intra + intra_vertical_indices = ( + vertical_topk[vertical_topk >= prev_chunk_end_pos] + - prev_chunk_end_pos + ) + if intra_vertical_indices.nelement() == 0: + intra_vertical_indices = torch.cat( + [ + intra_vertical_indices, + torch.arange( + 0, + k_states_intra.size(0), + max(1, k_states_intra.size(0) / 5), + dtype=torch.int32, + device=intra_vertical_indices.device, + ), + ] + ) + slash_topk = slash_topk_buffer[head_i, : heads_slash_size[head_i]] + intra_slash_indices = (qk.size(-1) - 1) - slash_topk[ + slash_topk >= prev_chunk_end_pos + ] + # fill buffer + v_count = intra_vertical_indices.nelement() + s_count = intra_slash_indices.nelement() + vertical_size_buffer[head_i] = v_count + slash_sizes_buffer[head_i] = s_count + vertical_buffer[head_i, :v_count].copy_(intra_vertical_indices) + slash_buffer[head_i, :s_count].copy_(intra_slash_indices) + # succ + if prev_chunk_end_pos - chunk_len >= 0: + succ_vertical_indices = vertical_topk[ + (vertical_topk < prev_chunk_end_pos) + & (vertical_topk >= prev_chunk_end_pos - chunk_len) + ] - (prev_chunk_end_pos - chunk_len) + # TODO: support no vertical + if succ_vertical_indices.nelement() == 0: + succ_vertical_indices = torch.cat( + [ + succ_vertical_indices, + torch.arange( + 0, + k_states_succ.size(0), + max(1, k_states_succ.size(0) / 5), + dtype=torch.int32, + device=intra_vertical_indices.device, + ), + ] + ) + succ_slash_indices = ( + prev_chunk_end_pos + (qend - qbegin) - 1 + ) - slash_topk[ + ( + (slash_topk >= (prev_chunk_end_pos - chunk_len)) + & (slash_topk < (prev_chunk_end_pos + (qend - qbegin))) + ) + ] + if succ_slash_indices.nelement() == 0: + succ_slash_indices = torch.cat( + [ + succ_slash_indices, + torch.arange( + 0, + k_states_succ.size(0), + max(1, k_states_succ.size(0) / 5), + dtype=torch.int32, + device=intra_vertical_indices.device, + ), + ] + ) + # fill buffer + v_count = succ_vertical_indices.nelement() + s_count = succ_slash_indices.nelement() + succ_vertical_size_buffer[head_i] = v_count + succ_slash_sizes_buffer[head_i] = s_count + succ_vertical_buffer[head_i, :v_count].copy_( + succ_vertical_indices + ) + succ_slash_buffer[head_i, :s_count].copy_(succ_slash_indices) + + if prev_chunk_end_pos - 2 * chunk_len >= 0: + inter_vertical_indices = vertical_topk[ + vertical_topk < prev_chunk_end_pos - chunk_len + ] + + if inter_vertical_indices.nelement() == 0: + inter_vertical_indices = torch.cat( + [ + inter_vertical_indices, + torch.arange( + 0, + k_states_inter.size(0), + max(1, k_states_inter.size(0) / 5), + dtype=torch.int32, + device=intra_vertical_indices.device, + ), + ] + ) + inter_slash_indices = ( + prev_chunk_end_pos - chunk_len + (qend - qbegin) - 1 + ) - slash_topk[ + slash_topk + < (prev_chunk_end_pos - chunk_len + (qend - qbegin)) + ] + if inter_slash_indices.nelement() == 0: + inter_slash_indices = torch.cat( + [ + inter_slash_indices, + torch.arange( + 0, + k_states_inter.size(0), + max(1, k_states_inter.size(0) / 5), + dtype=torch.int32, + device=intra_vertical_indices.device, + ), + ] + ) + # fill buffer + v_count = inter_vertical_indices.nelement() + s_count = inter_slash_indices.nelement() + inter_vertical_size_buffer[head_i] = v_count + inter_slash_sizes_buffer[head_i] = s_count + inter_vertical_buffer[head_i, :v_count].copy_( + inter_vertical_indices + ) + inter_slash_buffer[head_i, :s_count].copy_(inter_slash_indices) + else: + intra_vertical_indices, intra_slash_indices = None, None + succ_vertical_indices, succ_slash_indices = None, None + inter_vertical_indices, inter_slash_indices = None, None + + if sparse_attn_enabled: + flash_result = self._do_flash_attn( + q_states_intra, + k_states_intra, + v_states_intra, + softmax_scale=softmax_scale, + causal=True, + stage="intra", + vertical_indices=vertical_buffer, + slash_indices=slash_buffer, + vertical_indices_count=vertical_size_buffer, + slash_indices_count=slash_sizes_buffer, + mergehead_softmax_scale=softmax_scale, + sparse_attn_enabled=sparse_attn_enabled, + ) + else: + flash_result = self._do_flash_attn( + q_states_intra, + k_states_intra, + v_states_intra, + softmax_scale=softmax_scale, + causal=True, + stage="intra", + vertical_indices=intra_vertical_indices, + slash_indices=intra_slash_indices, + sparse_attn_enabled=sparse_attn_enabled, + ) + flash_per_chunk.append(flash_result) + + if prev_chunk_end_pos - chunk_len >= 0: + if sparse_attn_enabled: + flash_result = self._do_flash_attn( + q_states_succ, + k_states_succ, + v_states_succ, + softmax_scale=softmax_scale, + causal=False, + stage="succ", + vertical_indices=succ_vertical_buffer, + slash_indices=succ_slash_buffer, + vertical_indices_count=succ_vertical_size_buffer, + slash_indices_count=succ_slash_sizes_buffer, + mergehead_softmax_scale=softmax_scale, + sparse_attn_enabled=sparse_attn_enabled, + ) + else: + flash_result = self._do_flash_attn( + q_states_succ, + k_states_succ, + v_states_succ, + softmax_scale=softmax_scale, + causal=False, + stage="succ", + vertical_indices=succ_vertical_indices, + slash_indices=succ_slash_indices, + sparse_attn_enabled=sparse_attn_enabled, + ) + flash_per_chunk.append(flash_result) + + if prev_chunk_end_pos - chunk_len * 2 >= 0: + if sparse_attn_enabled: + flash_result = self._do_flash_attn( + q_states_inter, + k_states_inter, + v_states_inter, + softmax_scale=softmax_scale, + causal=False, + stage="inter", + vertical_indices=inter_vertical_buffer, + slash_indices=inter_slash_buffer, + vertical_indices_count=inter_vertical_size_buffer, + slash_indices_count=inter_slash_sizes_buffer, + mergehead_softmax_scale=softmax_scale, + sparse_attn_enabled=sparse_attn_enabled, + ) + else: + flash_result = self._do_flash_attn( + q_states_inter, + k_states_inter, + v_states_inter, + softmax_scale=softmax_scale, + causal=False, + stage="inter", + vertical_indices=inter_vertical_indices, + slash_indices=inter_slash_indices, + sparse_attn_enabled=sparse_attn_enabled, + ) + flash_per_chunk.append(flash_result) + + flash_results.append(flash_per_chunk) + begin = end + + attn_output = self._merge_attn_outputs(flash_results) + del flash_results + return attn_output + + def _do_flash_attn( + self, + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + softmax_scale: float, + causal: bool = True, + max_seqlen_k: Optional[int] = None, + stage: str = "intra", + vertical_indices: Optional[torch.Tensor] = None, + slash_indices: Optional[torch.Tensor] = None, + vertical_indices_count: Optional[torch.Tensor] = None, + slash_indices_count: Optional[torch.Tensor] = None, + mergehead_softmax_scale: Optional[float] = None, + sparse_attn_enabled: Optional[bool] = False, + ): + if max_seqlen_k is None: + max_seqlen_k = key_states.shape[0] + + q_len = query_states.shape[0] + q_heads = query_states.shape[1] + h_dim = query_states.shape[-1] + + if sparse_attn_enabled: + assert slash_indices is not None + if stage == "intra": + assert causal + else: + assert not causal + + query_states = query_states.unsqueeze(0).transpose(1, 2) + key_states = key_states.unsqueeze(0).transpose(1, 2) + value_states = value_states.unsqueeze(0).transpose(1, 2) + + q = query_states + k = key_states + v = value_states + + if vertical_indices_count is not None and slash_indices_count is not None: + assert mergehead_softmax_scale is not None + + res, s_lse = _vertical_slash_sparse_attention( + q, + k, + v, + vertical_indices, + slash_indices, + mergehead_softmax_scale, + causal=causal, + stage=stage, + vertical_indices_count=vertical_indices_count, + slash_indices_count=slash_indices_count, + ) + res = res.view(q_heads, q_len, h_dim).transpose( + 0, 1 + ) # (qlen,nhead,h_dim) + s_lse = ( + s_lse.view(q_heads, q_len, 1).squeeze(-1).unsqueeze(0).float() + ) # (1, nhead,qlen) + else: + res, s_lse = _vertical_slash_sparse_attention( + q, + k, + v, + vertical_indices, + slash_indices, + softmax_scale, + causal=causal, + stage=stage, + ) + res = res.view(q_len, q_heads, h_dim) + s_lse = s_lse.view(q_len, q_heads, 1).transpose(0, 2).float() + return res, s_lse + + output, softmax_lse, *rest = flash_attn_varlen_func( + q=query_states, + k=key_states, + v=value_states, + softmax_scale=softmax_scale, + cu_seqlens_q=torch.tensor( + [0, query_states.shape[0]], + dtype=torch.int32, + device=query_states.device, + ), + max_seqlen_q=query_states.shape[0], + cu_seqlens_k=torch.tensor( + [0, max_seqlen_k], dtype=torch.int32, device=query_states.device + ), + max_seqlen_k=max_seqlen_k, + causal=causal, + return_softmax_lse=True, + ) + softmax_lse = softmax_lse.view(q_len, q_heads, 1).transpose(0, 2).float() + return output, softmax_lse + + def _merge_attn_outputs( + self, + flash_results: List[List[Tuple[torch.Tensor, torch.Tensor]]], + return_lse: Optional[bool] = False, + ) -> torch.Tensor: + attn_outputs_all = [] + logits_all = [] + + for flash_per_chunk in flash_results: + if len(flash_per_chunk) == 1: + attn_outputs_all.append(flash_per_chunk[0][0]) + if return_lse: + logits_all.append(flash_per_chunk[0][1]) + continue + + attn_outputs = torch.stack( + [flash_attn_output[0] for flash_attn_output in flash_per_chunk] + ) + logits = torch.stack( + [flash_attn_output[1] for flash_attn_output in flash_per_chunk] + ) + logits = logits.to(torch.float32) + + if return_lse: + max_val = torch.max(logits, dim=0).values + diff = torch.abs(logits[0] - logits[1]) + log_sum_exp = max_val + torch.log1p(torch.exp(-diff)) + logits_all.append(log_sum_exp) + + max_logits = torch.max(logits, dim=0).values + stable_logits = logits - max_logits.unsqueeze(0) + lse_s = torch.exp(stable_logits).detach() + lse_sum = torch.sum(lse_s, dim=0) + lse_s /= lse_sum + attn_outputs *= lse_s.unsqueeze(-1).transpose(2, 3).squeeze(1) + attn_outputs_all.append(attn_outputs.sum(dim=0)) + + if return_lse: + return (torch.cat(attn_outputs_all, dim=0), torch.cat(logits_all, dim=-1)) + else: + return torch.cat(attn_outputs_all, dim=0) + + def _dual_chunk_flash_attn_decoding( + self, + query: torch.Tensor, + query_succ: torch.Tensor, + query_inter: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + softmax_scale: float, + causal: bool, + chunk_size: int, + local_size: int, + original_max_position_embeddings: int, + decode_meta: DualChunkFlashAttentionMetadata, + ): + if not causal: + raise ValueError("Dual Chunk Attention does not support causal=False") + + block_size = value_cache.shape[1] + chunk_len = chunk_size - local_size + if chunk_len % block_size != 0: + raise ValueError("chunk_len must be divisible by block_size.") + if original_max_position_embeddings > 0: + assert decode_meta.scaling_factor is not None + scaling_factor = decode_meta.scaling_factor + query = (query * scaling_factor.view(-1, 1, 1, 1)).to( + query.dtype + ) # possible for numerical issue, need to fused in the kernel + query_succ = (query_succ * scaling_factor.view(-1, 1, 1, 1)).to(query.dtype) + query_inter = (query_inter * scaling_factor.view(-1, 1, 1, 1)).to( + query.dtype + ) + outputs_list = [] + softmax_lses_list = [] + + # intra-attention + intra_output, intra_softmax_lse = ( + self._dual_chunk_flash_attn_decoding_with_exp_sums( + query, + key_cache, + value_cache, + decode_meta.block_tables_intra, + decode_meta.seq_lens_intra, + softmax_scale, + causal=False, + ) + ) + outputs_list.append(intra_output) + softmax_lses_list.append(intra_softmax_lse) + + # succ-attention + if decode_meta.max_seq_len_succ: + succ_output, succ_softmax_lse = ( + self._dual_chunk_flash_attn_decoding_with_exp_sums( + query_succ, + key_cache, + value_cache, + decode_meta.block_tables_succ, + decode_meta.seq_lens_succ, + softmax_scale, + causal=False, + ) + ) + outputs_list.append(succ_output) + softmax_lses_list.append(succ_softmax_lse) + + # inter-attention + if decode_meta.max_seq_len_inter: + inter_output, inter_softmax_lse = ( + self._dual_chunk_flash_attn_decoding_with_exp_sums( + query_inter, + key_cache, + value_cache, + block_table[:, : decode_meta.max_seq_len_inter], + decode_meta.seq_lens_inter, + softmax_scale, + causal=False, + ) + ) + outputs_list.append(inter_output) + softmax_lses_list.append(inter_softmax_lse) + outputs = torch.stack(outputs_list, dim=0) + del outputs_list + softmax_lses = torch.stack(softmax_lses_list, dim=0).to(torch.float32) + del softmax_lses_list + max_logits = torch.max(softmax_lses, dim=0).values + stable_logits = softmax_lses - max_logits.unsqueeze(0) + lse_s = torch.exp(stable_logits).detach() + lse_sum = torch.sum(lse_s, dim=0) + lse_s /= lse_sum + outputs *= lse_s.unsqueeze(-1).transpose(2, 3) + return outputs.sum(0) + + def _dual_chunk_flash_attn_decoding_with_exp_sums( + self, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + softmax_scale: float, + causal: bool, + ): + out, softmax_lse, *rest_expand = flash_attn_with_kvcache( + q=query, + k_cache=key_cache, + v_cache=value_cache, + page_table=block_table, + cache_seqlens=cache_seqlens, + softmax_scale=softmax_scale, + causal=causal, + return_softmax_lse=True, + ) + mask = cache_seqlens == 0 + out[mask] = 0 + softmax_lse[mask] = -float("inf") + return out, softmax_lse + + +def _vertical_slash_sparse_attention( + query: torch.Tensor, # [BATCH, N_HEADS, N_CTX, D_HEAD] + key: torch.Tensor, # [BATCH, N_HEADS, N_KV_CTX, D_HEAD] + value: torch.Tensor, # [BATCH, N_HEADS, N_KV_CTX, D_HEAD] + v_idx: torch.Tensor, # [BATCH, N_HEADS, NNZ_V] + s_idx: torch.Tensor, # [BATCH, N_HEADS, NNZ_S] + softmax_scale: float, + causal: bool = True, + stage: str = "intra", + block_size_M: int = 64, + block_size_N: int = 64, + vertical_indices_count: torch.Tensor = None, # [N_HEADS,] + slash_indices_count: torch.Tensor = None, +): + if stage == "intra": + assert causal + else: + assert not causal + + batch_size, num_heads, context_size, head_dim = query.shape + _, _, kv_seq_len, _ = key.shape + + if head_dim not in [16, 32, 64, 128, 256, 512]: + target_dim = 2 ** math.ceil(math.log2(head_dim)) - head_dim + query = F.pad(query, [0, target_dim, 0, 0, 0, 0, 0, 0]) + key = F.pad(key, [0, target_dim, 0, 0, 0, 0, 0, 0]) + value = F.pad(value, [0, target_dim, 0, 0, 0, 0, 0, 0]) + + v_idx = ( + v_idx.to(torch.int32) + .reshape((batch_size, num_heads, -1)) + .sort(dim=-1, descending=False)[0] + ) + s_idx = ( + s_idx.to(torch.int32) + .reshape((batch_size, num_heads, -1)) + .sort(dim=-1, descending=True)[0] + ) + q_seqlens = torch.tensor([context_size], dtype=torch.int32, device=query.device) + kv_seqlens = torch.tensor([kv_seq_len], dtype=torch.int32, device=query.device) + + if vertical_indices_count is not None and slash_indices_count is not None: + ( + block_count, + block_offset, + column_count, + column_index, + ) = convert_vertical_slash_indexes_mergehead( + q_seqlens, + kv_seqlens, + v_idx, + s_idx, + vertical_indices_count, + slash_indices_count, + context_size, + block_size_M, + block_size_N, + causal, + ) + else: + ( + block_count, + block_offset, + column_count, + column_index, + ) = convert_vertical_slash_indexes( + q_seqlens, + kv_seqlens, + v_idx, + s_idx, + context_size, + block_size_M, + block_size_N, + causal, + ) + + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + out, lse = sparse_attn_func( + q, + k, + v, + block_count, + block_offset, + column_count, + column_index, + causal=causal, + softmax_scale=softmax_scale, + return_softmax_lse=True, + ) + out = out.transpose(1, 2).contiguous() + softmax_lse = lse.reshape(*lse.shape, 1) + return (out[..., :context_size, :head_dim], softmax_lse[..., :context_size, :]) + + +def _sum_all_diagonal_matrix(mat: torch.tensor): + h, n, m = mat.shape + # Zero matrix used for padding + zero_mat = torch.zeros((h, n, n), device=mat.device) + # pads the matrix on left and right + mat_padded = torch.cat((zero_mat, mat, zero_mat), -1) + # Change the strides + mat_strided = mat_padded.as_strided( + (1, n, n + m), (n * (2 * n + m), 2 * n + m + 1, 1) + ) + # Sums the resulting matrix's columns + sum_diags = torch.sum(mat_strided, 1) + return sum_diags[:, 1:] # drop left bottom corner + + +def _get_block(block_table: torch.Tensor, block_size: int, begin: int, end: int): + begin_block = begin // block_size + end_block = (end - 1) // block_size + 1 + return block_table[begin_block:end_block] diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py new file mode 100644 index 000000000..a48a9e649 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/chunk.py @@ -0,0 +1,242 @@ +# Adapted from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/gated_delta_rule/chunk.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings +from typing import Optional + +import torch +from einops import rearrange + +from sglang.srt.layers.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from sglang.srt.layers.attention.fla.chunk_o import chunk_fwd_o +from sglang.srt.layers.attention.fla.chunk_scaled_dot_kkt import ( + chunk_scaled_dot_kkt_fwd, +) +from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum +from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd +from sglang.srt.layers.attention.fla.solve_tril import solve_tril +from sglang.srt.layers.attention.fla.utils import ( + SUPPRESS_LEVEL, + autocast_custom_fwd, + input_guard, +) +from sglang.srt.layers.attention.fla.wy_fast import recompute_w_u_fwd + + +def chunk_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: Optional[torch.LongTensor] = None, +): + g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens) + # obtain WY representation. u is actually the new v. + A = chunk_scaled_dot_kkt_fwd( + k=k, beta=beta, g_cumsum=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32 + ) + A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype) + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g_cumsum=g, + cu_seqlens=cu_seqlens, + ) + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + ) + if SUPPRESS_LEVEL < 3: + return g, o, A, final_state, None, None, None + elif SUPPRESS_LEVEL >= 3: + return g, o, A, final_state, w, h, v_new + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: Optional[torch.LongTensor] = None, + use_qk_l2norm_in_kernel: bool = False, + ): + q_orig = q + k_orig = k + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + + g, o, A, final_state, w, h, v_new = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o.to(q.dtype), final_state + + +@torch.compiler.disable +def chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: bool = False, + use_qk_l2norm_in_kernel: bool = False, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. + g (torch.Tensor): + (forget) gating tensor (in log space!) of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`. + beta (torch.Tensor): + betas of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`. + scale (Optional[int]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format, which is not supported for variable-length inputs. + Default: `False`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + assert q.dtype == k.dtype == v.dtype + assert ( + q.dtype != torch.float32 + ), "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16." + assert ( + len(beta.shape) == 3 + ), "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise." + + if head_first: + raise DeprecationWarning( + "head_first is deprecated and will be removed in a future version. " + "Please use head_first=False for now instead." + ) + q, k, v, beta, g = map( + lambda x: rearrange(x, "b h t ... -> b t h ..."), (q, k, v, beta, g) + ) + # if not head_first and q.shape[1] < q.shape[2]: + # warnings.warn( + # f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " + # "This may indicate the inputs were passed in head-first format [B, H, T, ...] " + # "when head_first=False was specified. " + # "Please verify your input tensor format matches the expected shape [B, T, H, ...]." + # ) + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkGatedDeltaRuleFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + use_qk_l2norm_in_kernel, + ) + if head_first: + o = rearrange(o, "b t h ... -> b h t ...") + return o, final_state diff --git a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py new file mode 100644 index 000000000..8f4109abd --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py @@ -0,0 +1,890 @@ +# Adapted from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/common/chunk_delta_h.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# from typing import Optional, Tuple + +# import torch +# import triton +# import triton.language as tl + +# from sglang.srt.layers.attention.fla.index import ( +# prepare_chunk_indices, +# prepare_chunk_offsets, +# ) +# from sglang.srt.layers.attention.fla.op import exp, safe_exp +# from sglang.srt.layers.attention.fla.utils import is_nvidia_hopper + +# NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16] + + +# @triton.heuristics( +# { +# "USE_G": lambda args: args["g"] is not None, +# "USE_INITIAL_STATE": lambda args: args["h0"] is not None, +# "STORE_FINAL_STATE": lambda args: args["ht"] is not None, +# "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, +# "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, +# } +# ) +# # @triton.autotune( +# # configs=[ +# # triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) +# # for num_warps in [2, 4] +# # for num_stages in [2, 3, 4] +# # for BV in [32, 64] +# # ], +# # key=["H", "K", "V", "BT", "USE_G"], +# # use_cuda_graph=use_cuda_graph, +# # ) +# @triton.jit(do_not_specialize=["T"]) +# def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( +# k, +# v, +# w, +# v_new, +# g, +# h, +# h0, +# ht, +# cu_seqlens, +# chunk_offsets, +# T, +# H: tl.constexpr, +# Hg: tl.constexpr, +# K: tl.constexpr, +# V: tl.constexpr, +# BT: tl.constexpr, +# BV: tl.constexpr, +# USE_G: tl.constexpr, +# USE_INITIAL_STATE: tl.constexpr, +# STORE_FINAL_STATE: tl.constexpr, +# SAVE_NEW_VALUE: tl.constexpr, +# IS_VARLEN: tl.constexpr, +# ): +# i_v, i_nh = tl.program_id(0), tl.program_id(1) +# i_n, i_h = i_nh // H, i_nh % H +# if IS_VARLEN: +# bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( +# cu_seqlens + i_n + 1 +# ).to(tl.int32) +# T = eos - bos +# NT = tl.cdiv(T, BT) +# boh = tl.load(chunk_offsets + i_n).to(tl.int32) +# else: +# bos, eos = i_n * T, i_n * T + T +# NT = tl.cdiv(T, BT) +# boh = i_n * NT + +# # [BK, BV] +# b_h1 = tl.zeros([64, BV], dtype=tl.float32) +# if K > 64: +# b_h2 = tl.zeros([64, BV], dtype=tl.float32) +# if K > 128: +# b_h3 = tl.zeros([64, BV], dtype=tl.float32) +# if K > 192: +# b_h4 = tl.zeros([64, BV], dtype=tl.float32) + +# # calculate offset +# h += (boh * H + i_h) * K * V +# v += (bos * H + i_h) * V +# k += (bos * Hg + i_h // (H // Hg)) * K +# w += (bos * H + i_h) * K +# if SAVE_NEW_VALUE: +# v_new += (bos * H + i_h) * V +# stride_v = H * V +# stride_h = H * K * V +# stride_k = Hg * K +# stride_w = H * K +# if USE_INITIAL_STATE: +# h0 = h0 + i_nh * K * V +# if STORE_FINAL_STATE: +# ht = ht + i_nh * K * V + +# # load initial state +# if USE_INITIAL_STATE: +# p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) +# b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) +# if K > 64: +# p_h0_2 = tl.make_block_ptr( +# h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) +# ) +# b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) +# if K > 128: +# p_h0_3 = tl.make_block_ptr( +# h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) +# ) +# b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) +# if K > 192: +# p_h0_4 = tl.make_block_ptr( +# h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) +# ) +# b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + +# # main recurrence +# for i_t in range(NT): +# p_h1 = tl.make_block_ptr( +# h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0) +# ) +# tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) +# if K > 64: +# p_h2 = tl.make_block_ptr( +# h + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) +# ) +# tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) +# if K > 128: +# p_h3 = tl.make_block_ptr( +# h + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) +# ) +# tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) +# if K > 192: +# p_h4 = tl.make_block_ptr( +# h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) +# ) +# tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + +# p_v = tl.make_block_ptr( +# v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) +# ) +# p_v_new = ( +# tl.make_block_ptr( +# v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) +# ) +# if SAVE_NEW_VALUE +# else None +# ) +# b_v_new = tl.zeros([BT, BV], dtype=tl.float32) +# p_w = tl.make_block_ptr( +# w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) +# ) +# b_w = tl.load(p_w, boundary_check=(0, 1)) +# b_v_new += tl.dot(b_w, b_h1.to(b_w.dtype)) +# if K > 64: +# p_w = tl.make_block_ptr( +# w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) +# ) +# b_w = tl.load(p_w, boundary_check=(0, 1)) +# b_v_new += tl.dot(b_w, b_h2.to(b_w.dtype)) +# if K > 128: +# p_w = tl.make_block_ptr( +# w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) +# ) +# b_w = tl.load(p_w, boundary_check=(0, 1)) +# b_v_new += tl.dot(b_w, b_h3.to(b_w.dtype)) +# if K > 192: +# p_w = tl.make_block_ptr( +# w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) +# ) +# b_w = tl.load(p_w, boundary_check=(0, 1)) +# b_v_new += tl.dot(b_w, b_h4.to(b_w.dtype)) +# b_v_new = -b_v_new + tl.load(p_v, boundary_check=(0, 1)) + +# if SAVE_NEW_VALUE: +# p_v_new = tl.make_block_ptr( +# v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) +# ) +# tl.store( +# p_v_new, b_v_new.to(p_v_new.dtype.element_ty), boundary_check=(0, 1) +# ) + +# if USE_G: +# last_idx = min((i_t + 1) * BT, T) - 1 +# b_g_last = tl.load(g + bos * H + last_idx * H + i_h) +# p_g = tl.make_block_ptr( +# g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) +# ) +# b_g = tl.load(p_g, boundary_check=(0,)) +# b_v_new = b_v_new * safe_exp(b_g_last - b_g)[:, None] +# b_g_last = exp(b_g_last) +# b_h1 = b_h1 * b_g_last +# if K > 64: +# b_h2 = b_h2 * b_g_last +# if K > 128: +# b_h3 = b_h3 * b_g_last +# if K > 192: +# b_h4 = b_h4 * b_g_last +# b_v_new = b_v_new.to(k.dtype.element_ty) +# p_k = tl.make_block_ptr( +# k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) +# ) +# b_k = tl.load(p_k, boundary_check=(0, 1)) +# b_h1 += tl.dot(b_k, b_v_new) +# if K > 64: +# p_k = tl.make_block_ptr( +# k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) +# ) +# b_k = tl.load(p_k, boundary_check=(0, 1)) +# b_h2 += tl.dot(b_k, b_v_new) +# if K > 128: +# p_k = tl.make_block_ptr( +# k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) +# ) +# b_k = tl.load(p_k, boundary_check=(0, 1)) +# b_h3 += tl.dot(b_k, b_v_new) +# if K > 192: +# p_k = tl.make_block_ptr( +# k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) +# ) +# b_k = tl.load(p_k, boundary_check=(0, 1)) +# b_h4 += tl.dot(b_k, b_v_new) + +# # epilogue +# if STORE_FINAL_STATE: +# p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) +# tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) +# if K > 64: +# p_ht = tl.make_block_ptr( +# ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) +# ) +# tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) +# if K > 128: +# p_ht = tl.make_block_ptr( +# ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) +# ) +# tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) +# if K > 192: +# p_ht = tl.make_block_ptr( +# ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) +# ) +# tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +# def chunk_gated_delta_rule_fwd_h( +# k: torch.Tensor, +# w: torch.Tensor, +# u: torch.Tensor, +# g: Optional[torch.Tensor] = None, +# initial_state: Optional[torch.Tensor] = None, +# output_final_state: bool = False, +# chunk_size: int = 64, # SY: remove this argument and force chunk size 64? +# save_new_value: bool = True, +# cu_seqlens: Optional[torch.LongTensor] = None, +# ) -> Tuple[torch.Tensor, torch.Tensor]: +# B, T, Hg, K, V = *k.shape, u.shape[-1] +# H = u.shape[-2] +# BT = chunk_size + +# chunk_indices = ( +# prepare_chunk_indices(cu_seqlens, chunk_size) +# if cu_seqlens is not None +# else None +# ) +# # N: the actual number of sequences in the batch with either equal or variable lengths +# if cu_seqlens is None: +# N, NT, chunk_offsets = B, triton.cdiv(T, BT), None +# else: +# N, NT, chunk_offsets = ( +# len(cu_seqlens) - 1, +# len(chunk_indices), +# prepare_chunk_offsets(cu_seqlens, BT), +# ) +# assert K <= 256, "current kernel does not support head dimension larger than 256." + +# h = k.new_empty(B, NT, H, K, V) +# final_state = ( +# k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None +# ) + +# v_new = torch.empty_like(u) if save_new_value else None + +# def grid(meta): +# return (triton.cdiv(V, meta["BV"]), N * H) +# # if hasattr(torch.version, 'hip') and torch.version.hip is not None: +# # # HIP环境回退实现 +# # return h, v_new, final_state +# chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( +# k=k, +# v=u, +# w=w, +# v_new=v_new, +# g=g, +# h=h, +# h0=initial_state, +# ht=final_state, +# cu_seqlens=cu_seqlens, +# chunk_offsets=chunk_offsets, +# T=T, +# H=H, +# Hg=Hg, +# K=K, +# V=V, +# BT=BT, +# BV=32, +# num_warps=4, +# num_stages=2, +# ) +# return h, v_new, final_state + + +from typing import Optional, Tuple +import math +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.index import ( + prepare_chunk_indices, + prepare_chunk_offsets, +) +from sglang.srt.layers.attention.fla.op import exp, safe_exp +from sglang.srt.layers.attention.fla.utils import is_nvidia_hopper + +# 保守一些在 HIP/ROCm 上更稳；NVIDIA 可调大 +NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16] + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, # [B, T, Hg, K] + v, # [B, T, H, V] + w, # [B, T, H, K] + v_new, # [B, T, H, V] or None + g, # [B, T, H] or None + h, # [B, NT, H, K, V] + h0, # [B, H, K, V] or None + ht, # [B, H, K, V] or None + cu_seqlens, # [N+1] or None + chunk_offsets,# [N] or None + T, # 全局 T（固定长时就是实际长度） + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, # chunk size（通常 64） + BV: tl.constexpr, # V 方向 tile 宽度（通常 32） + NT: tl.constexpr, # host 计算的上界（chunk 数） + USE_G: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # 程序实例：X 维为 V/BV，Y 维为 N*H + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n = i_nh // H + i_h = i_nh % H + + # 计算该序列 runtime 真实长度 TR、时间起点 bos、h 的 chunk 偏移 boh + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + TR = (eos - bos).to(tl.int32) + TR = tl.maximum(TR, 0) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos = (i_n * T).to(tl.int32) + eos = (i_n * T + T).to(tl.int32) + TR = T.to(tl.int32) + boh = (i_n * NT).to(tl.int32) + + # 预计算步长（以扁平指针计） + stride_h_t = H * K * V # h 时间块步长 + stride_k_hg = Hg * K # k 的 Hg*K 步长 + stride_w_hk = H * K # w 的 H*K 步长 + stride_v_hv = H * V # v/v_new 的 H*V 步长 + + # 移动到当前 sequence/head 的基址 + h = h + (boh * H + i_h) * (K * V) + v = v + (bos * stride_v_hv + i_h * V) + if SAVE_NEW_VALUE: + v_new = v_new + (bos * stride_v_hv + i_h * V) + k = k + (bos * stride_k_hg + (i_h // (H // Hg)) * K) + w = w + (bos * stride_w_hk + i_h * K) + + # hidden state tiles + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # 初始状态 + if USE_INITIAL_STATE: + h0 = h0 + i_nh * (K * V) + p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + if K > 64: + p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + if K > 128: + p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + if K > 192: + p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + + # 最终状态写回基址 + if STORE_FINAL_STATE: + ht = ht + i_nh * (K * V) + + # 主循环（上界 NT） + for i_t in range(NT): + # 先把当前 hidden state 写到 h（合法，因为 h 的时间维是 NT 上界） + p_h1 = tl.make_block_ptr(h + i_t * stride_h_t, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr(h + i_t * stride_h_t, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr(h + i_t * stride_h_t, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr(h + i_t * stride_h_t, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + # 这个 chunk 是否有效（t0 >= TR 则整个 tile 越界 -> 跳过所有 block_ptr 操作） + t0 = i_t * BT + active = t0 < TR + if active: + # v_new = v - w @ h + # --- w @ h --- + b_v_new = tl.zeros([BT, BV], dtype=tl.float32) + + # w 的 0..63 列 + p_w = tl.make_block_ptr(w, (TR, K), (stride_w_hk, 1), (t0, 0), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) # 不能使用 mask + b_v_new += tl.dot(b_w, b_h1.to(b_w.dtype)) + + if K > 64: + p_w = tl.make_block_ptr(w, (TR, K), (stride_w_hk, 1), (t0, 64), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h2.to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr(w, (TR, K), (stride_w_hk, 1), (t0, 128), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h3.to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr(w, (TR, K), (stride_w_hk, 1), (t0, 192), (BT, 64), (1, 0)) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v_new += tl.dot(b_w, b_h4.to(b_w.dtype)) + + # v - (w @ h) + p_v = tl.make_block_ptr(v, (TR, V), (stride_v_hv, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + b_v_chunk = tl.load(p_v, boundary_check=(0, 1)) # 不能使用 mask + b_v_new = -b_v_new + b_v_chunk + + # 保存 v_new + if SAVE_NEW_VALUE: + p_vn = tl.make_block_ptr(v_new, (TR, V), (stride_v_hv, 1), (t0, i_v * BV), (BT, BV), (1, 0)) + tl.store(p_vn, b_v_new.to(p_vn.dtype.element_ty), boundary_check=(0, 1)) + + # gating：对 g 用普通指针 + mask 读取（不是 block_ptr） + if USE_G: + offs_t = t0 + tl.arange(0, BT) # [BT] + t_mask = offs_t < TR # [BT] + g_base = g + bos * H + i_h + b_g = tl.load(g_base + offs_t * H, mask=t_mask, other=0.0) # ✔ 这里不是 block_ptr，可用 mask + + t_end = tl.minimum(t0 + BT, TR) + last_idx = tl.maximum(t_end - 1, 0).to(tl.int32) + b_g_last = tl.load(g_base + last_idx * H, mask=(TR > 0), other=0.0) + + gating_factor = safe_exp(b_g_last - b_g)[:, None] + b_v_new = b_v_new * gating_factor + + g_last_exp = exp(b_g_last) + b_h1 = b_h1 * g_last_exp + if K > 64: + b_h2 = b_h2 * g_last_exp + if K > 128: + b_h3 = b_h3 * g_last_exp + if K > 192: + b_h4 = b_h4 * g_last_exp + + # h += k^T @ v_new + b_v_new_cast = b_v_new.to(k.dtype.element_ty) + p_k = tl.make_block_ptr(k, (K, TR), (1, stride_k_hg), (0, t0), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) # 不能使用 mask + b_h1 += tl.dot(b_k, b_v_new_cast) + + if K > 64: + p_k = tl.make_block_ptr(k, (K, TR), (1, stride_k_hg), (64, t0), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h2 += tl.dot(b_k, b_v_new_cast) + if K > 128: + p_k = tl.make_block_ptr(k, (K, TR), (1, stride_k_hg), (128, t0), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h3 += tl.dot(b_k, b_v_new_cast) + if K > 192: + p_k = tl.make_block_ptr(k, (K, TR), (1, stride_k_hg), (192, t0), (64, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h4 += tl.dot(b_k, b_v_new_cast) + + # 末态写回 + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, # [B, T, Hg, K] + w: torch.Tensor, # [B, T, H, K] + u: torch.Tensor, # [B, T, H, V] + g: Optional[torch.Tensor] = None, # [B, T, H] + initial_state: Optional[torch.Tensor] = None, # [B, H, K, V] + output_final_state: bool = False, + chunk_size: int = 64, + save_new_value: bool = True, + cu_seqlens: Optional[torch.LongTensor] = None, # [N+1] 或 None +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """ + Host wrapper： + - 计算 NT（以 chunk 为单位的上界，constexpr）； + - kernel 内对块是否有效做 runtime 守卫（active）； + - block_ptr 的读写只用 boundary_check，严禁 mask/other； + - 对 g 使用标量指针 + mask 读取。 + """ + B, T, Hg, K = k.shape + H, V = u.shape[-2], u.shape[-1] + BT = chunk_size + + if cu_seqlens is None: + N = B + NT = math.ceil(T / BT) # 固定长：全局上界 + chunk_offsets = None + else: + N = len(cu_seqlens) - 1 + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) # 变长：所有序列的总 chunk 数 + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + + assert K <= 256, "current kernel does not support head dimension larger than 256." + + h = k.new_empty(B, NT, H, K, V) + final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None + v_new = torch.empty_like(u) if save_new_value else None + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), N * H) + + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + BV=32, + NT=NT, + num_warps=2, + num_stages=1, + ) + return h, v_new, final_state + + + + +# from typing import Optional, Tuple +# import math +# import torch +# import triton +# import triton.language as tl + +# from sglang.srt.layers.attention.fla.index import ( +# prepare_chunk_indices, +# prepare_chunk_offsets, +# ) +# from sglang.srt.layers.attention.fla.op import exp, safe_exp +# from sglang.srt.layers.attention.fla.utils import is_nvidia_hopper + +# # 保守一些在 HIP/ROCm 上更稳；NVIDIA 可调大 +# NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16] + + +# @triton.heuristics( +# { +# "USE_G": lambda args: args["g"] is not None, +# "USE_INITIAL_STATE": lambda args: args["h0"] is not None, +# "STORE_FINAL_STATE": lambda args: args["ht"] is not None, +# "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, +# "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, +# } +# ) +# @triton.jit(do_not_specialize=["T"]) +# def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( +# k, # [B, T, Hg, K] +# v, # [B, T, H, V] +# w, # [B, T, H, K] +# v_new, # [B, T, H, V] or None +# g, # [B, T, H] or None +# h, # [B, NT, H, K, V] +# h0, # [B, H, K, V] or None +# ht, # [B, H, K, V] or None +# cu_seqlens, # [N+1] or None +# chunk_offsets,# [N] or None +# T, # 全局 T（固定长时就是实际长度） +# H: tl.constexpr, +# Hg: tl.constexpr, +# K: tl.constexpr, +# V: tl.constexpr, +# BT: tl.constexpr, # chunk size（通常 64） +# BV: tl.constexpr, # V 方向 tile 宽度（通常 32） +# NT: tl.constexpr, # host 计算的上界（chunk 数） +# USE_G: tl.constexpr, +# USE_INITIAL_STATE: tl.constexpr, +# STORE_FINAL_STATE: tl.constexpr, +# SAVE_NEW_VALUE: tl.constexpr, +# IS_VARLEN: tl.constexpr, +# ): +# # 程序实例：X 维为 V/BV，Y 维为 N*H +# i_v, i_nh = tl.program_id(0), tl.program_id(1) +# i_n = i_nh // H +# i_h = i_nh % H + +# # 计算该序列 runtime 真实长度 TR、时间起点 bos、h 的 chunk 偏移 boh +# if IS_VARLEN: +# bos = tl.load(cu_seqlens + i_n).to(tl.int32) +# eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) +# TR = (eos - bos).to(tl.int32) +# TR = tl.maximum(TR, 0) +# boh = tl.load(chunk_offsets + i_n).to(tl.int32) +# else: +# bos = (i_n * T).to(tl.int32) +# eos = (i_n * T + T).to(tl.int32) +# TR = T.to(tl.int32) +# boh = (i_n * NT).to(tl.int32) + +# # 预计算步长（以扁平指针计） +# stride_h_t = H * K * V # h 时间块步长 +# stride_k_hg = Hg * K # k 的 Hg*K 步长 +# stride_w_hk = H * K # w 的 H*K 步长 +# stride_v_hv = H * V # v/v_new 的 H*V 步长 + +# # 移动到当前 sequence/head 的基址 +# h = h + (boh * H + i_h) * (K * V) +# v = v + (bos * stride_v_hv + i_h * V) +# if SAVE_NEW_VALUE: +# v_new = v_new + (bos * stride_v_hv + i_h * V) +# k = k + (bos * stride_k_hg + (i_h // (H // Hg)) * K) +# w = w + (bos * stride_w_hk + i_h * K) + +# # hidden state tiles +# b_h1 = tl.zeros([64, BV], dtype=tl.float32) +# if K > 64: +# b_h2 = tl.zeros([64, BV], dtype=tl.float32) +# if K > 128: +# b_h3 = tl.zeros([64, BV], dtype=tl.float32) +# if K > 192: +# b_h4 = tl.zeros([64, BV], dtype=tl.float32) + +# # 初始状态 +# if USE_INITIAL_STATE: +# h0 = h0 + i_nh * (K * V) +# p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) +# b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) +# if K > 64: +# p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) +# b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) +# if K > 128: +# p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) +# b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) +# if K > 192: +# p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) +# b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + +# # 最终状态写回基址 +# if STORE_FINAL_STATE: +# ht = ht + i_nh * (K * V) + +# # 主循环（上界 NT） +# for i_t in range(NT): +# # 先把当前 hidden state 写到 h（合法，因为 h 的时间维是 NT 上界） +# p_h1 = tl.make_block_ptr(h + i_t * stride_h_t, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) +# tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) +# if K > 64: +# p_h2 = tl.make_block_ptr(h + i_t * stride_h_t, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) +# tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) +# if K > 128: +# p_h3 = tl.make_block_ptr(h + i_t * stride_h_t, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) +# tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) +# if K > 192: +# p_h4 = tl.make_block_ptr(h + i_t * stride_h_t, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) +# tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + +# # 这个 chunk 是否有效（t0 >= TR 则整个 tile 越界 -> 跳过所有 block_ptr 操作） +# t0 = i_t * BT +# active = t0 < TR +# if active: +# # v_new = v - w @ h +# # --- w @ h --- +# b_v_new = tl.zeros([BT, BV], dtype=tl.float32) + +# # w 的 0..63 列 +# p_w = tl.make_block_ptr(w, (TR, K), (stride_w_hk, 1), (t0, 0), (BT, 64), (1, 0)) +# b_w = tl.load(p_w, boundary_check=(0, 1)) # 不能使用 mask +# b_v_new += tl.dot(b_w, b_h1.to(b_w.dtype)) + +# if K > 64: +# p_w = tl.make_block_ptr(w, (TR, K), (stride_w_hk, 1), (t0, 64), (BT, 64), (1, 0)) +# b_w = tl.load(p_w, boundary_check=(0, 1)) +# b_v_new += tl.dot(b_w, b_h2.to(b_w.dtype)) +# if K > 128: +# p_w = tl.make_block_ptr(w, (TR, K), (stride_w_hk, 1), (t0, 128), (BT, 64), (1, 0)) +# b_w = tl.load(p_w, boundary_check=(0, 1)) +# b_v_new += tl.dot(b_w, b_h3.to(b_w.dtype)) +# if K > 192: +# p_w = tl.make_block_ptr(w, (TR, K), (stride_w_hk, 1), (t0, 192), (BT, 64), (1, 0)) +# b_w = tl.load(p_w, boundary_check=(0, 1)) +# b_v_new += tl.dot(b_w, b_h4.to(b_w.dtype)) + +# # v - (w @ h) +# p_v = tl.make_block_ptr(v, (TR, V), (stride_v_hv, 1), (t0, i_v * BV), (BT, BV), (1, 0)) +# b_v_chunk = tl.load(p_v, boundary_check=(0, 1)) # 不能使用 mask +# b_v_new = -b_v_new + b_v_chunk + +# # 保存 v_new +# if SAVE_NEW_VALUE: +# p_vn = tl.make_block_ptr(v_new, (TR, V), (stride_v_hv, 1), (t0, i_v * BV), (BT, BV), (1, 0)) +# tl.store(p_vn, b_v_new.to(p_vn.dtype.element_ty), boundary_check=(0, 1)) + +# # gating：对 g 用普通指针 + mask 读取（不是 block_ptr） +# if USE_G: +# offs_t = t0 + tl.arange(0, BT) # [BT] +# t_mask = offs_t < TR # [BT] +# g_base = g + bos * H + i_h +# b_g = tl.load(g_base + offs_t * H, mask=t_mask, other=0.0) # ✔ 这里不是 block_ptr，可用 mask + +# t_end = tl.minimum(t0 + BT, TR) +# last_idx = tl.maximum(t_end - 1, 0).to(tl.int32) +# b_g_last = tl.load(g_base + last_idx * H, mask=(TR > 0), other=0.0) + +# gating_factor = safe_exp(b_g_last - b_g)[:, None] +# b_v_new = b_v_new * gating_factor + +# g_last_exp = exp(b_g_last) +# b_h1 = b_h1 * g_last_exp +# if K > 64: +# b_h2 = b_h2 * g_last_exp +# if K > 128: +# b_h3 = b_h3 * g_last_exp +# if K > 192: +# b_h4 = b_h4 * g_last_exp + +# # h += k^T @ v_new +# b_v_new_cast = b_v_new.to(k.dtype.element_ty) +# p_k = tl.make_block_ptr(k, (K, TR), (1, stride_k_hg), (0, t0), (64, BT), (0, 1)) +# b_k = tl.load(p_k, boundary_check=(0, 1)) # 不能使用 mask +# b_h1 += tl.dot(b_k, b_v_new_cast) + +# if K > 64: +# p_k = tl.make_block_ptr(k, (K, TR), (1, stride_k_hg), (64, t0), (64, BT), (0, 1)) +# b_k = tl.load(p_k, boundary_check=(0, 1)) +# b_h2 += tl.dot(b_k, b_v_new_cast) +# if K > 128: +# p_k = tl.make_block_ptr(k, (K, TR), (1, stride_k_hg), (128, t0), (64, BT), (0, 1)) +# b_k = tl.load(p_k, boundary_check=(0, 1)) +# b_h3 += tl.dot(b_k, b_v_new_cast) +# if K > 192: +# p_k = tl.make_block_ptr(k, (K, TR), (1, stride_k_hg), (192, t0), (64, BT), (0, 1)) +# b_k = tl.load(p_k, boundary_check=(0, 1)) +# b_h4 += tl.dot(b_k, b_v_new_cast) + +# # 末态写回 +# if STORE_FINAL_STATE: +# p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) +# tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) +# if K > 64: +# p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) +# tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) +# if K > 128: +# p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) +# tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) +# if K > 192: +# p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) +# tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +# def chunk_gated_delta_rule_fwd_h( +# k: torch.Tensor, # [B, T, Hg, K] +# w: torch.Tensor, # [B, T, H, K] +# u: torch.Tensor, # [B, T, H, V] +# g: Optional[torch.Tensor] = None, # [B, T, H] +# initial_state: Optional[torch.Tensor] = None, # [B, H, K, V] +# output_final_state: bool = False, +# chunk_size: int = 64, +# save_new_value: bool = True, +# cu_seqlens: Optional[torch.LongTensor] = None, # [N+1] 或 None +# ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: +# """ +# Host wrapper： +# - 计算 NT（以 chunk 为单位的上界，constexpr）； +# - kernel 内对块是否有效做 runtime 守卫（active）； +# - block_ptr 的读写只用 boundary_check，严禁 mask/other； +# - 对 g 使用标量指针 + mask 读取。 +# """ +# B, T, Hg, K = k.shape +# H, V = u.shape[-2], u.shape[-1] +# BT = chunk_size + +# if cu_seqlens is None: +# N = B +# NT = math.ceil(T / BT) # 固定长：全局上界 +# chunk_offsets = None +# else: +# N = len(cu_seqlens) - 1 +# chunk_indices = prepare_chunk_indices(cu_seqlens, BT) +# NT = len(chunk_indices) # 变长：所有序列的总 chunk 数 +# chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) + +# assert K <= 256, "current kernel does not support head dimension larger than 256." + +# h = k.new_empty(B, NT, H, K, V) +# final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None +# v_new = torch.empty_like(u) if save_new_value else None + +# def grid(meta): +# return (triton.cdiv(V, meta["BV"]), N * H) + +# chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( +# k=k, +# v=u, +# w=w, +# v_new=v_new, +# g=g, +# h=h, +# h0=initial_state, +# ht=final_state, +# cu_seqlens=cu_seqlens, +# chunk_offsets=chunk_offsets, +# T=T, +# H=H, +# Hg=Hg, +# K=K, +# V=V, +# BT=BT, +# BV=32, +# NT=NT, +# num_warps=2, +# num_stages=1, +# ) +# return h, v_new, final_state diff --git a/python/sglang/srt/layers/attention/fla/chunk_o.py b/python/sglang/srt/layers/attention/fla/chunk_o.py new file mode 100644 index 000000000..d672c646b --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/chunk_o.py @@ -0,0 +1,178 @@ +# Adapted from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/common/chunk_o.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.index import prepare_chunk_indices +from sglang.srt.layers.attention.fla.op import exp, safe_exp +from sglang.srt.layers.attention.fla.utils import check_shared_mem, is_nvidia_hopper + +BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] +NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8] + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +# @triton.autotune( +# configs=[ +# triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) +# for BK in BKV_LIST +# for BV in BKV_LIST +# for num_warps in NUM_WARPS +# for num_stages in [2, 3, 4] +# ], +# key=["H", "K", "V", "BT"], +# ) +@triton.jit(do_not_specialize=["T"]) +def chunk_fwd_kernel_o( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * Hg + i_h // (H // Hg)) * K + k += (bos * Hg + i_h // (H // Hg)) * K + v += (bos * H + i_h) * V + o += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K * V + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) + ) + p_h = tl.make_block_ptr( + h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) + ) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + + # [BT, BK] @ [BK, BV] -> [BT, BV] + b_o += tl.dot(b_q, b_h) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_A += tl.dot(b_q, b_k) + + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_o = b_o * exp(b_g)[:, None] + b_A = b_A * safe_exp(b_g[:, None] - b_g[None, :]) + + o_i = tl.arange(0, BT) + m_A = o_i[:, None] >= o_i[None, :] + b_A = tl.where(m_A, b_A, 0) + + p_v = tl.make_block_ptr( + v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + p_o = tl.make_block_ptr( + o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # to fix mma -> mma layout conversion + # already solved by triton v3.2 or higher + b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_fwd_o( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: Optional[torch.Tensor] = None, # cumsum of log decay + scale: Optional[float] = None, + cu_seqlens: Optional[torch.LongTensor] = None, + chunk_size: int = 64, +) -> torch.Tensor: + B, T, Hg, K, V = *q.shape, v.shape[-1] + H = v.shape[-2] + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + if scale is None: + scale = k.shape[-1] ** -0.5 + + o = torch.empty_like(v) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), NT, B * H) + + chunk_fwd_kernel_o[grid]( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T=T, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + BK=128, + BV=64, + num_warps=4, + num_stages=2, + ) + return o diff --git a/python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py b/python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py new file mode 100644 index 000000000..699350d31 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py @@ -0,0 +1,151 @@ +# Adapted from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/common/chunk_scaled_dot_kkt.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.index import prepare_chunk_indices +from sglang.srt.layers.attention.fla.op import safe_exp + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + "USE_G": lambda args: args["g_cumsum"] is not None, + } +) +# @triton.autotune( +# configs=[ +# triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) +# for BK in [32, 64, 128] +# for num_warps in [2, 4, 8] +# for num_stages in [2, 3, 4] +# ], +# key=["H", "K", "BT", "IS_VARLEN"], +# ) +@triton.jit(do_not_specialize=["T"]) +def chunk_scaled_dot_kkt_fwd_kernel( + k, + beta, + g_cumsum, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = tl.arange(0, BT) + + p_beta = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * Hg + i_h // (H // Hg)) * K, + (T, K), + (Hg * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_beta[:, None] + b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k)) + + if USE_G: + p_g = tl.make_block_ptr( + g_cumsum + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_diff = b_g[:, None] - b_g[None, :] + b_A = b_A * safe_exp(b_g_diff) + + b_A = tl.where(o_t[:, None] > o_t[None, :], b_A, 0) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_scaled_dot_kkt_fwd( + k: torch.Tensor, + beta: torch.Tensor, + g_cumsum: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, + chunk_size: int = 64, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + r""" + Compute beta * K * K^T. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + g_cumsum (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H]`. + Default: None + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + + B, T, Hg, K = k.shape + + H = beta.shape[-1] + BT = chunk_size + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) + chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)]( + k=k, + beta=beta, + g_cumsum=g_cumsum, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + Hg=Hg, + K=K, + BT=BT, + BK=64, + num_warps=8, + num_stages=3, + ) + return A diff --git a/python/sglang/srt/layers/attention/fla/cumsum.py b/python/sglang/srt/layers/attention/fla/cumsum.py new file mode 100644 index 000000000..b8e3cdde1 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/cumsum.py @@ -0,0 +1,300 @@ +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/utils/cumsum.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.index import prepare_chunk_indices +from sglang.srt.layers.attention.fla.utils import check_shared_mem, input_guard + +BS_LIST = [32, 64] if check_shared_mem() else [16, 32] + + +@triton.heuristics( + { + "HAS_SCALE": lambda args: args["scale"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +# @triton.autotune( +# configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], +# key=["B", "H", "BT", "IS_VARLEN", "REVERSE"], +# ) +@triton.jit(do_not_specialize=["T"]) +def chunk_local_cumsum_scalar_kernel( + s, + o, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if HEAD_FIRST: + p_s = tl.make_block_ptr( + s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) + ) + p_o = tl.make_block_ptr( + o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) + ) + else: + p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + # [BT] + b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_o = tl.cumsum(b_s, axis=0) + if REVERSE: + b_z = tl.sum(b_s, axis=0) + b_o = -b_o + b_z[None] + b_s + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics( + { + "HAS_SCALE": lambda args: args["scale"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BS": BS}, num_warps=num_warps) + for BS in BS_LIST + for num_warps in [2, 4, 8] + ], + key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_local_cumsum_vector_kernel( + s, + o, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BT) + if REVERSE: + m_s = tl.where(o_i[:, None] <= o_i[None, :], 1.0, 0.0) + else: + m_s = tl.where(o_i[:, None] >= o_i[None, :], 1.0, 0.0) + + if HEAD_FIRST: + p_s = tl.make_block_ptr( + s + (bos * H + i_h * T) * S, + (T, S), + (S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h * T) * S, + (T, S), + (S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + else: + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_o = tl.dot(m_s, b_s, allow_tf32=False) + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_local_cumsum_scalar( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + scale: float = None, + cu_seqlens: Optional[torch.Tensor] = None, + head_first: bool = False, + output_dtype: Optional[torch.dtype] = torch.float, +) -> torch.Tensor: + if head_first: + B, H, T = g.shape + else: + B, T, H = g.shape + assert chunk_size == 2 ** ( + chunk_size.bit_length() - 1 + ), "chunk_size must be a power of 2" + BT = chunk_size + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + grid = (NT, B * H) + chunk_local_cumsum_scalar_kernel[grid]( + s=g_org, + o=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + B=B, + H=H, + BT=BT, + HEAD_FIRST=head_first, + REVERSE=reverse, + num_warps=8, + num_stages=3, + ) + return g + + +def chunk_local_cumsum_vector( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + scale: float = None, + cu_seqlens: Optional[torch.Tensor] = None, + head_first: bool = False, + output_dtype: Optional[torch.dtype] = torch.float, +) -> torch.Tensor: + if head_first: + B, H, T, S = g.shape + else: + B, T, H, S = g.shape + BT = chunk_size + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + assert chunk_size == 2 ** ( + chunk_size.bit_length() - 1 + ), "chunk_size must be a power of 2" + + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + + def grid(meta): + return (triton.cdiv(meta["S"], meta["BS"]), NT, B * H) + + # keep cumulative normalizer in fp32 + # this kernel is equivalent to + # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1) + chunk_local_cumsum_vector_kernel[grid]( + s=g_org, + o=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + B=B, + H=H, + S=S, + BT=BT, + HEAD_FIRST=head_first, + REVERSE=reverse, + ) + return g + + +@input_guard +def chunk_local_cumsum( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + scale: float = None, + cu_seqlens: Optional[torch.Tensor] = None, + head_first: bool = False, + output_dtype: Optional[torch.dtype] = torch.float, + **kwargs, +) -> torch.Tensor: + if cu_seqlens is not None: + assert ( + g.shape[0] == 1 + ), "Only batch size 1 is supported when cu_seqlens are provided" + if len(g.shape) == 3: + return chunk_local_cumsum_scalar( + g=g, + chunk_size=chunk_size, + reverse=reverse, + scale=scale, + cu_seqlens=cu_seqlens, + head_first=head_first, + output_dtype=output_dtype, + ) + elif len(g.shape) == 4: + return chunk_local_cumsum_vector( + g=g, + chunk_size=chunk_size, + reverse=reverse, + scale=scale, + cu_seqlens=cu_seqlens, + head_first=head_first, + output_dtype=output_dtype, + ) + else: + raise ValueError( + f"Unsupported input shape {g.shape}, " + f"which should be (B, T, H, D) if `head_first=False` " + f"or (B, H, T, D) otherwise" + ) diff --git a/python/sglang/srt/layers/attention/fla/fused_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_recurrent.py new file mode 100644 index 000000000..fa7262ce2 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/fused_recurrent.py @@ -0,0 +1,640 @@ +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/gated_delta_rule/fused_recurrent.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.op import exp +from sglang.srt.layers.attention.fla.utils import input_guard + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.jit(do_not_specialize=["T"]) +def fused_recurrent_gated_delta_rule_fwd_kernel( + q, + k, + v, + g, + beta, + o, + h0, + ht, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + STORE_FINAL_STATE: tl.constexpr, # whether to store final state + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + (bos * HV + i_hv) * V + o_v + else: + p_beta = beta + bos * HV + i_hv + p_g = g + bos * HV + i_hv + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_g = tl.load(p_g).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6) + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6) + b_q = b_q * scale + # [BK, BV] + b_h *= exp(b_g) + # [BV] + b_v -= tl.sum(b_h * b_k[:, None], 0) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_v *= b_beta + # [BK, BV] + b_h += b_k[:, None] * b_v[None, :] + # [BV] + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + p_g += HV + p_beta += HV * (V if IS_BETA_HEADWISE else 1) + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 3 + num_warps = 1 + + o = q.new_empty(NK, *v.shape) + if output_final_state: + final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) + else: + final_state = None + + grid = (NK, NV, N * HV) + fused_recurrent_gated_delta_rule_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim == v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o, final_state + + +class FusedRecurrentFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: Optional[torch.LongTensor] = None, + use_qk_l2norm_in_kernel: bool = False, + ): + o, final_state = fused_recurrent_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + cu_seqlens=cu_seqlens, + ) + + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht): + raise NotImplementedError( + "Backward pass is not implemented yet and we do not have plans to implement it " + "because we haven't figured out how to compute dg without materializing the full " + "hidden states for all time steps." + ) + + +def fused_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None, + use_qk_l2norm_in_kernel: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA is applied if `HV > H`. + g (torch.Tensor): + g (decays) of shape `[B, T, HV]`. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[int]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, device='cuda') + >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) + >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() + >>> h0 = torch.randn(B, HV, K, V, device='cuda') + >>> o, ht = fused_gated_recurrent_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_gated_recurrent_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "scale must be positive" + if beta is None: + beta = torch.ones_like(q[..., 0]) + o, final_state = FusedRecurrentFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + use_qk_l2norm_in_kernel, + ) + return o, final_state + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["h0_source"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + "CACHE_INTERMEDIATE_STATES": lambda args: args["intermediate_states_buffer"] + is not None, + } +) +@triton.jit(do_not_specialize=["T"]) +def fused_recurrent_gated_delta_rule_update_fwd_kernel( + q, + k, + v, + g, + beta, + o, + h0_source, + h0_indices, + cu_seqlens, + scale, + intermediate_states_buffer, + cache_steps, + T, + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, + DISABLE_STATE_UPDATE: tl.constexpr, # whether to disable final state update + DISABLE_OUTPUT_CALCULATION: tl.constexpr, # whether to disable output calculation + CACHE_INTERMEDIATE_STATES: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + (bos * HV + i_hv) * V + o_v + else: + p_beta = beta + bos * HV + i_hv + p_g = g + bos * HV + i_hv + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n) + # Add bounds checking for idx + if idx >= 0: # Assuming negative indices are invalid + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + # Prepare intermediate state cache variables if enabled + cache_idx = -1 + if CACHE_INTERMEDIATE_STATES: + cache_idx = tl.load(h0_indices + i_n) + + step_idx = 0 + for _ in range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_g = tl.load(p_g).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6) + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6) + b_q = b_q * scale + # [BK, BV] + b_h *= exp(b_g) + # [BV] + b_v -= tl.sum(b_h * b_k[:, None], 0) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_v *= b_beta + # [BK, BV] + b_h += b_k[:, None] * b_v[None, :] + # [BV] + if not DISABLE_OUTPUT_CALCULATION: + b_o = tl.sum(b_h * b_q[:, None], 0) + # core attn output + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + # store intermediate states if enabled + if CACHE_INTERMEDIATE_STATES: + if cache_idx >= 0: + # Compute cache pointer for this step + step_offset = step_idx * HV * K * V + cache_ptr = ( + intermediate_states_buffer + + cache_idx * cache_steps * HV * K * V + + step_offset + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h) + + step_idx += 1 + + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + p_g += HV + p_beta += HV * (V if IS_BETA_HEADWISE else 1) + + # Store final state back to h0_source with bounds checking + # ssm states + if not DISABLE_STATE_UPDATE: + idx = tl.load(h0_indices + i_n) + if idx >= 0: # Add bounds checking + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_gated_delta_rule_update_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None, + disable_state_update: bool = False, + disable_output_calculation: bool = False, + intermediate_states_buffer: Optional[torch.Tensor] = None, + cache_steps: Optional[int] = None, +) -> torch.Tensor: + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 3 + num_warps = 1 + + if disable_output_calculation: + # When output calculation is disabled, allocate minimal tensor + o = q.new_empty(NK, 1, 1, 1, 1) # minimal allocation + else: + o = q.new_empty(NK, *v.shape) + + grid = (NK, NV, N * HV) + + fused_recurrent_gated_delta_rule_update_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + o=o, + h0_source=initial_state_source, + h0_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + scale=scale, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=0 if cache_steps is None else cache_steps, + T=T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim == v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + DISABLE_STATE_UPDATE=disable_state_update, + DISABLE_OUTPUT_CALCULATION=disable_output_calculation, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o + + +class FusedRecurrentUpdateFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + cu_seqlens: Optional[torch.LongTensor] = None, + use_qk_l2norm_in_kernel: bool = False, + disable_state_update: bool = False, + disable_output_calculation: bool = False, + intermediate_states_buffer: Optional[torch.Tensor] = None, + cache_steps: Optional[int] = None, + ): + o = fused_recurrent_gated_delta_rule_update_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + cu_seqlens=cu_seqlens, + disable_state_update=disable_state_update, + disable_output_calculation=disable_output_calculation, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=cache_steps, + ) + + return o + + @staticmethod + @input_guard + def backward(ctx, do, dht): + raise NotImplementedError( + "Backward pass is not implemented yet and we do not have plans to implement it " + "because we haven't figured out how to compute dg without materializing the full " + "hidden states for all time steps." + ) + + +def fused_recurrent_gated_delta_rule_update( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor = None, + scale: float = None, + initial_state_source: torch.Tensor = None, + initial_state_indices: torch.Tensor = None, + cu_seqlens: Optional[torch.LongTensor] = None, + use_qk_l2norm_in_kernel: bool = False, + disable_state_update: bool = False, + disable_output_calculation: bool = False, + intermediate_states_buffer: Optional[torch.Tensor] = None, + cache_steps: Optional[int] = None, +) -> torch.Tensor: + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if ( + initial_state_source is not None + and initial_state_indices.shape[0] != len(cu_seqlens) - 1 + ): + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state_indices.shape[0]}." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "scale must be positive" + if beta is None: + beta = torch.ones_like(q[..., 0]) + o = FusedRecurrentUpdateFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state_source, + initial_state_indices, + cu_seqlens, + use_qk_l2norm_in_kernel, + disable_state_update, + disable_output_calculation, + intermediate_states_buffer, + cache_steps, + ) + return o diff --git a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py new file mode 100644 index 000000000..41837b980 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py @@ -0,0 +1,232 @@ +from typing import Optional + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.utils import input_guard + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["h0_source"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.jit(do_not_specialize=["T"]) +def fused_sigmoid_gating_delta_rule_update_kernel( + A_log, + a, + dt_bias, + softplus_beta, + softplus_threshold, + q, + k, + v, + b, + o, + h0_source, + h0_indices, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Fused kernel that combines sigmoid gating computation with recurrent delta rule update. + """ + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + p_b = b + bos * HV + i_hv + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + # Gating computation pointers + p_A_log = A_log + i_hv + p_a = a + bos * HV + i_hv + p_dt_bias = dt_bias + i_hv + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n) + if idx >= 0: + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + # Load inputs + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_b = tl.load(p_b).to(tl.float32) + + # Compute sigmoid gating + # Load gating parameters + b_A_log = tl.load(p_A_log).to(tl.float32) + b_a = tl.load(p_a).to(tl.float32) + b_dt_bias = tl.load(p_dt_bias).to(tl.float32) + + # Compute g = -exp(A_log) * softplus(a + dt_bias) + x = b_a + b_dt_bias + beta_x = softplus_beta * x + # Apply softplus with numerical stability + softplus_x = tl.where( + beta_x <= softplus_threshold, + (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), + x, + ) + b_g = -tl.exp(b_A_log) * softplus_x + + # Compute beta = sigmoid(b) + b_beta = 1.0 / (1.0 + tl.exp(-b_b)) + + # Apply L2 normalization if enabled + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6) + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6) + + b_q = b_q * scale + + # Apply gating to hidden state: h *= exp(g) + b_h *= tl.exp(b_g) + + # Delta rule: v -= sum(h * k, dim=0) + b_v -= tl.sum(b_h * b_k[:, None], 0) + + # Apply beta gating: v *= beta + b_v *= b_beta + + # Update hidden state: h += k[:, None] * v[None, :] + b_h += b_k[:, None] * b_v[None, :] + + # Compute output: o = sum(h * q, dim=0) + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + # Update pointers for next timestep + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + p_b += HV + p_a += HV + + # Store final state back to h0_source with bounds checking + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n) + if idx >= 0: + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + + +@input_guard +def fused_sigmoid_gating_delta_rule_update( + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + softplus_beta: float, + softplus_threshold: float, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, +): + """ + Fused triton implementation of sigmoid gating delta rule update. + This function uses a single fused kernel that combines both sigmoid gating computation + and the recurrent delta rule update for better performance. + """ + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 3 + num_warps = 1 + + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "scale must be positive" + + o = q.new_empty(NK, *v.shape) + grid = (NK, NV, N * HV) + + fused_sigmoid_gating_delta_rule_update_kernel[grid]( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + q=q, + k=k, + v=v, + b=b, + o=o, + h0_source=initial_state_source, + h0_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o diff --git a/python/sglang/srt/layers/attention/fla/index.py b/python/sglang/srt/layers/attention/fla/index.py new file mode 100644 index 000000000..754b98714 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/index.py @@ -0,0 +1,37 @@ +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/utils/index.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.utils import tensor_cache + + +@tensor_cache +def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: + return cu_seqlens[1:] - cu_seqlens[:-1] + + +@tensor_cache +def prepare_chunk_indices( + cu_seqlens: torch.LongTensor, chunk_size: int +) -> torch.LongTensor: + indices = torch.cat( + [ + torch.arange(n) + for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() + ] + ) + return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) + + +@tensor_cache +def prepare_chunk_offsets( + cu_seqlens: torch.LongTensor, chunk_size: int +) -> torch.LongTensor: + return torch.cat( + [cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)] + ).cumsum(-1) diff --git a/python/sglang/srt/layers/attention/fla/l2norm.py b/python/sglang/srt/layers/attention/fla/l2norm.py new file mode 100644 index 000000000..d6b6ae7f7 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/l2norm.py @@ -0,0 +1,150 @@ +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/modules/l2norm.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +import torch.nn as nn +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.utils import input_guard + +BT_LIST = [8, 16, 32, 64, 128] + + +# @triton.autotune( +# configs=[ +# triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16, 32] +# ], +# key=["D"], +# ) +@triton.jit +def l2norm_fwd_kernel1( + x, + y, + D, + BD: tl.constexpr, + eps, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + # Compute mean and variance + cols = tl.arange(0, BD) + mask = cols < D + b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) + b_var = tl.sum(b_x * b_x, axis=0) + b_rstd = 1 / tl.sqrt(b_var + eps) + # tl.store(Rstd + i_t, rstd) + # Normalize and apply linear transformation + b_y = b_x * b_rstd + tl.store(y + cols, b_y, mask=mask) + + +# @triton.autotune( +# configs=[ +# triton.Config({"BT": BT}, num_warps=num_warps) +# for num_warps in [1, 2, 4, 8, 16] +# for BT in BT_LIST +# ], +# key=["D", "NB"], +# ) +@triton.jit +def l2norm_fwd_kernel( + x, + y, + eps, + NB: tl.constexpr, + T: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, +): + i_t = tl.program_id(0) + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + b_var = tl.sum(b_x * b_x, axis=1) + b_y = b_x / tl.sqrt(b_var + eps)[:, None] + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +def l2norm_fwd( + x: torch.Tensor, eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None +): + x_shape_og = x.shape + x = x.view(-1, x.shape[-1]) + # allocate output + if output_dtype is None: + y = torch.empty_like(x) + else: + y = torch.empty_like(x, dtype=output_dtype) + assert y.stride(-1) == 1 + T, D = x.shape[0], x.shape[-1] + # rstd = torch.empty((T,), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer doesn't support feature dim >= 64KB.") + + if D <= 512: + NB = triton.cdiv(T, 2048) + + def grid(meta): + return (triton.cdiv(T, meta["BT"]),) + + l2norm_fwd_kernel[grid]( + x, + y, + eps, + NB=NB, + T=T, + D=D, + BD=BD, + BT=16, + num_warps=8, + num_stages=3, + ) + else: + l2norm_fwd_kernel1[(T,)]( + x, + y, + eps=eps, + D=D, + BD=BD, + num_warps=8, + num_stages=3, + ) + + return y.view(x_shape_og) + + +class L2NormFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward(ctx, x, eps=1e-6, output_dtype=None): + return l2norm_fwd(x, eps, output_dtype) + + +def l2norm( + x: torch.Tensor, eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None +) -> torch.Tensor: + return L2NormFunction.apply(x, eps, output_dtype) + + +l2_norm = l2norm + + +class L2Norm(nn.Module): + + def __init__(self, eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None): + super().__init__() + self.eps = eps + self.output_dtype = output_dtype + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return l2norm(x, self.eps, self.output_dtype) diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py new file mode 100644 index 000000000..bd53d0d64 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/layernorm_gated.py @@ -0,0 +1,326 @@ +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/modules/layernorm_gated.py +# Copyright (c) 2024, Tri Dao. +# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. +# This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling. +# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. + +import math + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange + + +def rms_norm_ref( + x, + weight, + bias, + z=None, + eps=1e-6, + group_size=None, + norm_before_gate=True, + upcast=True, +): + dtype = x.dtype + N = x.shape[-1] + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + z = z.float() if z is not None else z + if z is not None and not norm_before_gate: + x = x * F.silu(z) + if group_size is None: + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + else: + x_group = rearrange(x, "... (g d) -> ... g d", d=group_size) + rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps) + out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight + if bias is not None: + out = out + bias + if z is not None and norm_before_gate: + out *= F.silu(z) + return out.to(dtype) + + +@triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["Z"] is not None}) +@triton.jit +def _layer_norm_fwd_1pass_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_z_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_N: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + group = tl.program_id(1) + X += row * stride_x_row + group * N + Y += row * stride_y_row + group * N + if HAS_Z: + Z += row * stride_z_row + group * N + if not IS_RMS_NORM: + Mean += group * M + Rstd += group * M + W += group * N + if HAS_BIAS: + B += group * N + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + if HAS_Z and not NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=cols < N).to(tl.float32) + x *= z * tl.sigmoid(z) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean + row, mean) + xbar = tl.where(cols < N, x - mean, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + w = tl.load(W + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + y = x_hat * w + b if HAS_BIAS else x_hat * w + if HAS_Z and NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=mask).to(tl.float32) + y *= z * tl.sigmoid(z) + # Write output + tl.store(Y + cols, y, mask=mask) + + +def _layer_norm_fwd( + x, + weight, + bias, + eps, + z=None, + out=None, + group_size=None, + norm_before_gate=True, + is_rms_norm=False, +): + M, N = x.shape + if group_size is None: + group_size = N + assert N % group_size == 0 + ngroups = N // group_size + assert x.stride(-1) == 1 + if z is not None: + assert z.stride(-1) == 1 + assert z.shape == (M, N) + assert weight.shape == (N,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (N,) + # allocate output + if out is not None: + assert out.shape == x.shape + else: + out = torch.empty_like(x) + assert out.stride(-1) == 1 + mean = ( + torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) + if not is_rms_norm + else None + ) + rstd = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) + if group_size > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + num_warps = min(max(BLOCK_N // 256, 1), 8) + grid = (M, ngroups) + with torch.cuda.device(x.device.index): + _layer_norm_fwd_1pass_kernel[grid]( + x, + out, + weight, + bias, + z, + mean, + rstd, + x.stride(0), + out.stride(0), + z.stride(0) if z is not None else 0, + M, + group_size, + eps, + BLOCK_N=BLOCK_N, + NORM_BEFORE_GATE=norm_before_gate, + IS_RMS_NORM=is_rms_norm, + num_warps=num_warps, + ) + return out, mean, rstd + + +class LayerNormFn(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + x, + weight, + bias, + z=None, + eps=1e-6, + group_size=None, + norm_before_gate=True, + is_rms_norm=False, + ): + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" + + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + if x.stride(-1) != 1: + x = x.contiguous() + if z is not None: + assert z.shape == x_shape_og + z = z.reshape(-1, z.shape[-1]) + if z.stride(-1) != 1: + z = z.contiguous() + weight = weight.contiguous() + if bias is not None: + bias = bias.contiguous() + y, mean, rstd = _layer_norm_fwd( + x, + weight, + bias, + eps, + z=z, + group_size=group_size, + norm_before_gate=norm_before_gate, + is_rms_norm=is_rms_norm, + ) + return y.reshape(x_shape_og) + + +def layernorm_fn( + x, + weight, + bias, + z=None, + eps=1e-6, + group_size=None, + norm_before_gate=True, + is_rms_norm=False, +): + return LayerNormFn.apply( + x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm + ) + + +def rmsnorm_fn( + x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True +): + return LayerNormFn.apply( + x, weight, bias, z, eps, group_size, norm_before_gate, True + ) + + +class LayerNorm(torch.nn.Module): + + def __init__( + self, + hidden_size, + eps=1e-5, + group_size=None, + norm_before_gate=True, + device=None, + dtype=None, + ): + """If group_size is not None, we do GroupNorm with each group having group_size elements. + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). + """ + + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.group_size = group_size + self.norm_before_gate = norm_before_gate + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + torch.nn.init.zeros_(self.bias) + + def forward(self, x, z=None): + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" + return layernorm_fn( + x, + self.weight, + self.bias, + z=z, + group_size=self.group_size, + eps=self.eps, + norm_before_gate=self.norm_before_gate, + ) + + +class RMSNorm(torch.nn.Module): + + def __init__( + self, + hidden_size, + eps=1e-5, + group_size=None, + norm_before_gate=True, + device=None, + dtype=None, + ): + """If group_size is not None, we do GroupNorm with each group having group_size elements. + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.register_parameter("bias", None) + self.group_size = group_size + self.norm_before_gate = norm_before_gate + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + + def forward(self, x, z=None): + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" + return rmsnorm_fn( + x, + self.weight, + self.bias, + z=z, + eps=self.eps, + group_size=self.group_size, + norm_before_gate=self.norm_before_gate, + ) diff --git a/python/sglang/srt/layers/attention/fla/op.py b/python/sglang/srt/layers/attention/fla/op.py new file mode 100644 index 000000000..9b3191075 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/op.py @@ -0,0 +1,66 @@ +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/utils/op.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import os + +import triton +import triton.language as tl +import triton.language.extra.libdevice as tldevice + +from sglang.srt.layers.attention.fla.utils import is_gather_supported + +if os.environ.get("FLA_USE_FAST_OPS", "0") == "1": + exp = tldevice.fast_expf + exp2 = tldevice.exp2 + log = tldevice.fast_logf + log2 = tldevice.fast_log2f +else: + exp = tl.exp + exp2 = tl.math.exp2 + log = tl.log + log2 = tl.log2 + + +@triton.jit +def safe_exp(x): + return exp(tl.where(x <= 0, x, float("-inf"))) + + +if not is_gather_supported: + + @triton.jit + def gather(src, index, axis, _builder=None): + """ + Gather operation that works when tl.gather is not supported. + This is a fallback implementation that returns None. + Just to make triton compiler happy. + """ + return None + +else: + gather = tl.gather + + +if hasattr(triton.language, "_experimental_make_tensor_descriptor"): + # For Triton 3.3.x + make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor +elif hasattr(triton.language, "make_tensor_descriptor"): + # For Triton 3.4.x and later + make_tensor_descriptor = triton.language.make_tensor_descriptor +else: + """ + Fallback implementation when TMA is not supported. + Returns None to indicate TMA descriptors are unavailable. + Just make triton compiler happy. + """ + + @triton.jit + def make_tensor_descriptor( + base, + shape, + strides, + block_shape, + _builder=None, + ): + return None diff --git a/python/sglang/srt/layers/attention/fla/solve_tril.py b/python/sglang/srt/layers/attention/fla/solve_tril.py new file mode 100644 index 000000000..5c519507d --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/solve_tril.py @@ -0,0 +1,465 @@ +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/utils/solve_tril.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.index import prepare_chunk_indices +from sglang.srt.layers.attention.fla.utils import input_guard + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +# @triton.autotune( +# configs=[ +# triton.Config({}, num_warps=num_warps, num_stages=num_stages) +# for num_warps in [1, 2, 4, 8] +# for num_stages in [2, 3, 4, 5] +# ], +# key=["BT"], +# ) +@triton.jit(do_not_specialize=["T"]) +def solve_tril_16x16_kernel( + A, + Ad, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + A = A + (bos * H + i_h) * BT + Ad = Ad + (bos * H + i_h) * 16 + + offset = (i_t * 16) % BT + p_A = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0) + ) + p_Ai = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_A = -tl.where(tl.arange(0, 16)[:, None] > tl.arange(0, 16)[None, :], b_A, 0) + + o_i = tl.arange(0, 16) + for i in range(1, min(16, T - i_t * 16)): + b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset) + b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) + mask = o_i == i + b_A = tl.where(mask[:, None], b_a, b_A) + b_A += o_i[:, None] == o_i[None, :] + tl.store( + p_Ai, + b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +# @triton.autotune( +# configs=[ +# triton.Config({}, num_warps=num_warps, num_stages=num_stages) +# for num_warps in [1, 2, 4, 8] +# for num_stages in [2, 3, 4, 5] +# ], +# key=["H", "BT", "IS_VARLEN"], +# ) +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_32x32_inverse_kernel( + A, + Ad, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + A += (bos * H + i_h) * 32 + Ad += (bos * H + i_h) * 16 + Ai += (bos * H + i_h) * 32 + + p_A_21 = tl.make_block_ptr( + A, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0) + ) + p_Ad_11 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 32, 0), (16, 16), (1, 0) + ) + p_Ad_22 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0) + ) + p_Ai_11 = tl.make_block_ptr( + Ai, (T, 32), (H * 32, 1), (i_t * 32, 0), (16, 16), (1, 0) + ) + p_Ai_22 = tl.make_block_ptr( + Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 16), (16, 16), (1, 0) + ) + p_Ai_21 = tl.make_block_ptr( + Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0) + ) + + A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32) + Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32) + Ai_21 = -tl.dot( + tl.dot(Ai_22, A_21, input_precision="ieee"), Ai_11, input_precision="ieee" + ) + tl.store( + p_Ai_11, + Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +# @triton.autotune( +# configs=[ +# triton.Config({}, num_warps=num_warps, num_stages=num_stages) +# for num_warps in [2, 4, 8] +# for num_stages in [2, 3, 4, 5] +# ], +# key=["H", "BT", "IS_VARLEN"], +# ) +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_64x64_inverse_kernel( + A, + Ad, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + A += (bos * H + i_h) * 64 + Ad += (bos * H + i_h) * 16 + Ai += (bos * H + i_h) * 64 + + p_A_21 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0) + ) + p_A_32 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 32, 16), (16, 16), (1, 0) + ) + p_A_31 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0) + ) + p_A_43 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 32), (16, 16), (1, 0) + ) + p_A_42 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 16), (16, 16), (1, 0) + ) + p_A_41 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0) + ) + p_Ad_11 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 64, 0), (16, 16), (1, 0) + ) + p_Ad_22 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0) + ) + p_Ad_33 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0) + ) + p_Ad_44 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0) + ) + + A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) + A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) + A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) + A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) + + Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32) + Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32) + Ai_33 = tl.load(p_Ad_33, boundary_check=(0, 1)).to(tl.float32) + Ai_44 = tl.load(p_Ad_44, boundary_check=(0, 1)).to(tl.float32) + + Ai_21 = -tl.dot( + tl.dot(Ai_22, A_21, input_precision="ieee"), Ai_11, input_precision="ieee" + ) + Ai_32 = -tl.dot( + tl.dot(Ai_33, A_32, input_precision="ieee"), Ai_22, input_precision="ieee" + ) + Ai_43 = -tl.dot( + tl.dot(Ai_44, A_43, input_precision="ieee"), Ai_33, input_precision="ieee" + ) + + Ai_31 = -tl.dot( + Ai_33, + tl.dot(A_31, Ai_11, input_precision="ieee") + + tl.dot(A_32, Ai_21, input_precision="ieee"), + input_precision="ieee", + ) + Ai_42 = -tl.dot( + Ai_44, + tl.dot(A_42, Ai_22, input_precision="ieee") + + tl.dot(A_43, Ai_32, input_precision="ieee"), + input_precision="ieee", + ) + Ai_41 = -tl.dot( + Ai_44, + tl.dot(A_41, Ai_11, input_precision="ieee") + + tl.dot(A_42, Ai_21, input_precision="ieee") + + tl.dot(A_43, Ai_31, input_precision="ieee"), + input_precision="ieee", + ) + + p_Ai_11 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64, 0), (16, 16), (1, 0) + ) + p_Ai_22 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 16), (16, 16), (1, 0) + ) + p_Ai_33 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 32), (16, 16), (1, 0) + ) + p_Ai_44 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 48), (16, 16), (1, 0) + ) + p_Ai_21 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0) + ) + p_Ai_31 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0) + ) + p_Ai_32 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 16), (16, 16), (1, 0) + ) + p_Ai_41 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0) + ) + p_Ai_42 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 16), (16, 16), (1, 0) + ) + p_Ai_43 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 32), (16, 16), (1, 0) + ) + tl.store( + p_Ai_11, + Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_33, + Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_44, + Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_31, + Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_32, + Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_41, + Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_42, + Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_43, + Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + + fill_zeros = tl.zeros((16, 16), dtype=tl.float32) + p_Ai_12 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64, 16), (16, 16), (1, 0) + ) + p_Ai_13 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64, 32), (16, 16), (1, 0) + ) + p_Ai_14 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64, 48), (16, 16), (1, 0) + ) + p_Ai_23 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 32), (16, 16), (1, 0) + ) + p_Ai_24 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 48), (16, 16), (1, 0) + ) + p_Ai_34 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 48), (16, 16), (1, 0) + ) + tl.store( + p_Ai_12, + fill_zeros.to(p_Ai_12.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_13, + fill_zeros.to(p_Ai_13.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_14, + fill_zeros.to(p_Ai_14.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_23, + fill_zeros.to(p_Ai_23.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_24, + fill_zeros.to(p_Ai_24.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_34, + fill_zeros.to(p_Ai_34.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + + +@input_guard +def solve_tril( + A: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + output_dtype: torch.dtype = torch.float, +) -> torch.Tensor: + """ + Compute the inverse of the lower triangular matrix + A should be strictly lower triangular, i.e., A.triu() == 0. + + Args: + A (torch.Tensor): + [B, T, H, K] + cu_seqlens (torch.Tensor): + The cumulative sequence lengths of the input tensor. + Default: None. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float` + + Returns: + (I + A)^-1 with the same shape as A + """ + assert A.shape[-1] in [16, 32, 64] + + B, T, H, BT = A.shape + Ad = torch.empty( + B, T, H, 16, device=A.device, dtype=torch.float if BT != 16 else output_dtype + ) + + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, 16) if cu_seqlens is not None else None + ) + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, 16) + solve_tril_16x16_kernel[NT, B * H]( + A=A, + Ad=Ad, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + num_warps=1, + num_stages=4, + ) + if BT == 16: + return Ad + + Ai = torch.empty(B, T, H, BT, device=A.device, dtype=output_dtype) + merge_fn = ( + merge_16x16_to_32x32_inverse_kernel + if BT == 32 + else merge_16x16_to_64x64_inverse_kernel + ) + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) + merge_fn[NT, B * H]( + A=A, + Ad=Ad, + Ai=Ai, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + num_warps=4, + num_stages=3, + ) + return Ai diff --git a/python/sglang/srt/layers/attention/fla/utils.py b/python/sglang/srt/layers/attention/fla/utils.py new file mode 100644 index 000000000..3caf70de5 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/utils.py @@ -0,0 +1,331 @@ +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/utils.py +# -*- coding: utf-8 -*- + +import contextlib +import functools +import logging +import os +import sys +from enum import Enum +from functools import lru_cache +from typing import Any, Callable, Dict, Literal, Optional, Tuple + +import torch +import triton +from packaging import version + +logger = logging.getLogger(__name__) + +COMPILER_MODE = os.getenv("FLA_COMPILER_MODE") == "1" +FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1" + + +@lru_cache(maxsize=1) +def check_environments(): + """ + Checks the current operating system, Triton version, and Python version, + issuing warnings if they don't meet recommendations. + This function's body only runs once due to lru_cache. + """ + # Check Operating System + if sys.platform == "win32": + logger.warning( + "Detected Windows operating system. Triton does not have an official Windows release, " + "thus FLA will not be adapted for Windows, and any potential errors will not be fixed. " + "Please consider using a Linux environment for compatibility." + ) + + triton_version = version.parse(triton.__version__) + required_triton_version = version.parse("3.2.0") + + if triton_version < required_triton_version: + logger.warning( + f"Current Triton version {triton_version} is below the recommended 3.2.0 version. " + "Errors may occur and these issues will not be fixed. " + "Please consider upgrading Triton." + ) + + # Check Python version + py_version = version.parse(f"{sys.version_info.major}.{sys.version_info.minor}") + required_py_version = version.parse("3.11") + + if py_version < required_py_version: + logger.warning( + f"Current Python version {py_version} is below the recommended 3.11 version. " + "It is recommended to upgrade to Python 3.11 or higher for the best experience." + ) + + return None + + +check_environments() + + +def get_abs_err(x, y): + return (x.detach() - y.detach()).flatten().abs().max().item() + + +def get_err_ratio(x, y): + err = (x.detach() - y.detach()).flatten().square().mean().sqrt().item() + base = (x.detach()).flatten().square().mean().sqrt().item() + return err / (base + 1e-8) + + +def assert_close(prefix, ref, tri, ratio, warning=False, err_atol=1e-6): + abs_atol = get_abs_err(ref, tri) + msg = f"{prefix} diff: {abs_atol:.6f} ratio: {get_err_ratio(ref, tri):.6f}" + logger.info(msg) + error_rate = get_err_ratio(ref, tri) + if abs_atol <= err_atol: + return + if warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3)): + if error_rate > ratio: + import warnings + + warnings.warn(msg) + else: + assert error_rate < ratio, msg + + +SUPPRESS_LEVEL = int(os.getenv("GDN_RECOMPUTE_SUPPRESS_LEVEL", "0")) + + +def tensor_cache(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + """ + A decorator that caches the most recent results of a function with tensor inputs. + This decorator will store the output of the decorated function for the most recent set of input tensors. + The cache is limited to a fixed size (default is 4). When the cache is full, the oldest entry will be removed. + Args: + fn (Callable[..., torch.Tensor]): + The function to be decorated. It should take tensor inputs and return tensor outputs. + Returns: + Callable[..., torch.Tensor]: + A wrapped version of the input function with single-entry caching. + """ + + cache_entries: Tuple[Optional[Tuple], Optional[Dict], Any] = [] + cache_size = 4 + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + nonlocal cache_entries, cache_size + for i, entry in enumerate(cache_entries): + last_args, last_kwargs, last_result = entry + if len(args) == len(last_args) and len(kwargs) == len(last_kwargs): + if all(a is b for a, b in zip(args, last_args)) and all( + k in last_kwargs and v is last_kwargs[k] for k, v in kwargs.items() + ): + cache_entries = ( + cache_entries[:i] + + cache_entries[i + 1 :] + + [(args, kwargs, last_result)] + ) + return last_result + + result = fn(*args, **kwargs) + + if len(cache_entries) >= cache_size: + cache_entries = cache_entries[1:] + cache_entries.append((args, kwargs, result)) + return result + + return wrapper + + +def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + """ + A decorator to make sure all input tensors are contiguous and set the device based on input tensors. + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + contiguous_args = ( + i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args + ) + contiguous_kwargs = { + k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) + for k, v in kwargs.items() + } + + tensor = None + for arg in args: + if isinstance(arg, torch.Tensor): + tensor = arg + break + if tensor is None: + for value in kwargs.values(): + if isinstance(value, torch.Tensor): + tensor = value + break + + if tensor is not None: + ctx = custom_device_ctx(tensor.device.index) + else: + ctx = contextlib.nullcontext() + + with ctx: + return fn(*contiguous_args, **contiguous_kwargs) + + return wrapper + + +contiguous = input_guard + + +def require_version(version, hint): + """ + Perform a runtime check of the dependency versions, using the exact same syntax used by pip. + """ + + def decorator(fn): + @functools.wraps(fn) + def wrapper(ctx, *args, **kwargs): + from transformers.utils.versions import require_version + + require_version(version, hint) + return fn( + ctx, + *( + i if not isinstance(i, torch.Tensor) else i.contiguous() + for i in args + ), + **{ + k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) + for k, v in kwargs.items() + }, + ) + + return wrapper + + return decorator + + +def checkpoint(fn): + def wrapper(*args, **kwargs): + return torch.utils.checkpoint.checkpoint(fn, *args, **kwargs) + + return wrapper + + +@lru_cache(maxsize=None) +def check_pytorch_version(version_s: str = "2.4") -> bool: + return version.parse(torch.__version__) >= version.parse(version_s) + + +def _cpu_device_warning(): + import warnings + + warnings.warn( + ("Triton is not supported on current platform, roll back to CPU."), stacklevel=1 + ) + + +@lru_cache(maxsize=None) +def get_multiprocessor_count(tensor_idx: int = 0) -> int: + try: + return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)[ + "multiprocessor_count" + ] + except BaseException: + _cpu_device_warning() + return -1 + + +@lru_cache(maxsize=None) +def get_available_device() -> str: + try: + return triton.runtime.driver.active.get_current_target().backend + except BaseException: + _cpu_device_warning() + return "cpu" + + +@lru_cache(maxsize=None) +def _check_platform() -> Literal["nvidia", "amd", "intel", "musa"]: + device = get_available_device() + if device == "cuda": + return "nvidia" + elif device == "hip": + return "amd" + elif device == "xpu": + return "intel" + else: + return device + + +# For AMD GPUs, the triton backend is 'hip', while for Nvidia GPUs, the triton backend is 'cuda'. +# However, the torch backend is 'cuda' for both Nvidia and AMD GPUs. +# Therefore, we need to check the triton backend to determine the actual GPU vendor. +device = get_available_device() if get_available_device() != "hip" else "cuda" +device_torch_lib = getattr(torch, device) +device_platform = _check_platform() + +is_amd = device_platform == "amd" +is_intel = device_platform == "intel" +is_nvidia = device_platform == "nvidia" +is_intel_alchemist = is_intel and "Intel(R) Arc(TM) A" in torch.xpu.get_device_name(0) +is_nvidia_hopper = is_nvidia and ( + "NVIDIA H" in torch.cuda.get_device_name(0) + or torch.cuda.get_device_capability()[0] >= 9 +) +use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1" + +# Nvidia Ampere or newer, haven't check AMD and intel yet. +is_tf32_supported = is_nvidia and torch.cuda.get_device_capability(0)[0] >= 8 +is_gather_supported = hasattr(triton.language, "gather") + + +def get_all_max_shared_mem(): + try: + return [ + triton.runtime.driver.active.utils.get_device_properties(i)[ + "max_shared_mem" + ] + for i in range(device_torch_lib.device_count()) + ] + except BaseException: + _cpu_device_warning() + return [-1] + + +class Backend(Enum): + ADA = 101376 # RTX 4090 + AMPERE = 166912 # A100 + HOPPER = 232448 # H100 + DEFAULT = 102400 # Default + + @classmethod + def get_shared_memory(cls, arch: str) -> int: + try: + return cls[arch.upper()].value + except KeyError: + return cls.DEFAULT.value + + +@lru_cache(maxsize=None) +def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: + try: + device_shared_mem_list = get_all_max_shared_mem() + max_shared_memory = device_shared_mem_list[tensor_idx] + return max_shared_memory >= Backend.get_shared_memory(arch) + except Exception: + return False + + +if check_pytorch_version("2.4"): + device = "cuda" if device == "cpu" else device + autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device) + autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device) + + def custom_device_ctx(index: int): + return device_torch_lib.device(index) + +else: + assert ( + device == "cuda" + ), "Only cuda device is supported for PyTorch version < 2.4.0." + autocast_custom_fwd = device_torch_lib.amp.custom_fwd + autocast_custom_bwd = device_torch_lib.amp.custom_bwd + + def custom_device_ctx(index: int): + return torch.cuda.device(index) diff --git a/python/sglang/srt/layers/attention/fla/wy_fast.py b/python/sglang/srt/layers/attention/fla/wy_fast.py new file mode 100644 index 000000000..d51500eb4 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/wy_fast.py @@ -0,0 +1,158 @@ +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/gated_delta_rule/wy_fast.py +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.fla.index import prepare_chunk_indices +from sglang.srt.layers.attention.fla.op import safe_exp +from sglang.srt.layers.attention.fla.utils import check_shared_mem + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +# @triton.autotune( +# configs=[ +# triton.Config({}, num_warps=num_warps, num_stages=num_stages) +# for num_warps in [2, 4, 8] +# for num_stages in [2, 3, 4] +# ], +# key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +# ) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_beta = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_beta = tl.load(p_beta, boundary_check=(0,)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * Hg + i_h // (H // Hg)) * K, + (T, K), + (Hg * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) + b_w = tl.dot(b_A, b_kb) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g_cumsum: torch.Tensor, + A: torch.Tensor, + cu_seqlens: Optional[torch.LongTensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + B, T, Hg, K, V = *k.shape, v.shape[-1] + H = v.shape[-2] + BT = A.shape[-1] + + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BK = 64 + BV = 64 + u = torch.empty_like(v) + w = k.new_empty(B, T, H, K) + recompute_w_u_fwd_kernel[(NT, B * H)]( + k=k, + v=v, + beta=beta, + w=w, + u=u, + A=A, + g=g_cumsum, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + num_warps=4, + num_stages=3, + ) + return w, u + + +fwd_recompute_w_u = recompute_w_u_fwd diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py new file mode 100644 index 000000000..f7ca5e203 --- /dev/null +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -0,0 +1,2365 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Union + +import numpy as np +import torch +import triton +import triton.language as tl + +from sglang.srt.configs.model_config import AttentionArch +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.mem_cache.memory_pool import SWAKVPool +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + +from sgl_kernel import merge_state_v2 +from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache + + +@dataclass +class FlashAttentionMetadata: + """Metadata to be init once in the model forward pass, + each layer's forward pass can reuse the metadata. + + For each init metadata function, we will try set up them in below order + """ + + # Sequence lengths for the forward batch + cache_seqlens_int32: torch.Tensor = None + # Maximum sequence length for query + max_seq_len_q: int = 1 + # Maximum sequence length for key + max_seq_len_k: int = 0 + # Cumulative sequence lengths for query + cu_seqlens_q: torch.Tensor = None + # Cumulative sequence lengths for key + cu_seqlens_k: torch.Tensor = None + # Window size (typically used by Gemma) + window_size: tuple = (-1, -1) + # Page table, the index of KV Cache Tables/Blocks + page_table: torch.Tensor = None + + # Encoder metadata + # Cumulative sequence lengths for encoder key + encoder_cu_seqlens_k: torch.Tensor = None + # Maximum sequence length for encoder key + encoder_max_seq_len_k: int = 0 + # Sequence lengths for the forward batch + encoder_lens_int32: torch.Tensor = None + # Page table for the encoder + encoder_page_table: torch.Tensor = None + + @dataclass + class LocalAttentionMetadata: + local_query_start_loc: torch.Tensor = None # cu_seqlens_q for local attention + local_seqused_k: torch.Tensor = None # sequence lengths for local attention + local_block_table: torch.Tensor = None # block table for local attention + local_max_query_len: int = 0 # max query length for local attention + local_max_seq_len: int = 0 # max sequence length for local attention + + local_attn_metadata: Optional[LocalAttentionMetadata] = None + + # For sliding window attention topk>1 spec decoding + swa_spec_metadata: Optional[FlashAttentionMetadata] = None + + +# Copied from: +# https://github.com/houseroad/vllm/blob/4e45bfcaf928bdb9bd952b4ac922a3c205589ae8/vllm/v1/attention/backends/flash_attn.py +# +# Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into +# local attention blocks, where each block is passed to the attention kernel +# as an independent local ("virtual") batch item. +# +# For example, if are performing a chunked prefill a batch of 3 sequences: +# q_seqlens = [4, 10, 5] +# kv_seqlens = [6, 17, 9] +# Then normally for regular attention we would compute with an attention mask +# for batch idx 0 (q_seqlens = 4, kv_seqlens = 6) like: +# batch idx: 0 (q_seqlens = 4, kv_seqlens = 6) +# k_toks > 0 1 2 3 4 5 +# q_toks v _____________ +# 0 | 1 1 1 +# 1 | 1 1 1 1 +# 2 | 1 1 1 1 1 +# 3 | 1 1 1 1 1 1 +# +# for local attention (with attn_chunk_size = 4) we would compute with an +# attention mask like: +# batch idx: 0 (q_seqlens = 4, kv_seqlens = 6, attn_chunk_size = 4) +# k_toks > 0 1 2 3 4 5 +# q_toks v _____________ +# 0 | 1 1 1 +# 1 | 1 1 1 1 +# 2 | 1 +# 3 | 1 1 +# +# We can simulate this mask using standard flash-attention by breaking the +# sequences into local ("virtual") batches, where each local batch item is a +# local attention block, so in this case batch idx 0 would be broken up into: +# +# local-batch idx: 0 (q_seqlens = 2, kv_seqlens = 4) (batch 0) +# k_toks > 0 1 2 3 +# q_toks v _____________ +# 0 | 1 1 1 +# 1 | 1 1 1 1 +# local-batch idx: 1 (q_seqlens = 2, kv_seqlens = 2) (batch 0) +# k_toks > 4 5 +# q_toks v _____________ +# 2 | 1 +# 3 | 1 1 +# +# e.g. if we have: +# attn_chunk_size = 4 +# query_start_loc_np = [0, 4, 14, 19] (q_seqlens = [4, 10, 5]) +# Then this function would return: +# __b0__ ______b1______ __b2__ < orig batch indices +# q_seqlens_local = [ 2, 2, 1, 4, 4, 1, 4, 1] +# cu_seqlens_q_local = [0, 4, 6, 10, 14, 18, 19, 23, 24] +# seqlens_k_local = [ 4, 2, 4, 4, 4, 1, 4, 1] +# block_table_local : shape[local_virtual_batches, pages_per_local_batch] +def make_local_attention_virtual_batches( + attn_chunk_size: int, + query_start_loc_np: np.ndarray, + seq_lens_np: np.ndarray, + block_table: torch.Tensor, + page_size: int = 0, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, torch.Tensor]: + """ + Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into + local attention blocks, where each block is passed to the attention kernel + as an independent local ("virtual") batch item. + + Args: + attn_chunk_size: Size of local attention chunks + query_start_loc_np: Cumulative sum of query lengths (numpy array) + seq_lens_np: Sequence lengths (numpy array) + block_table: Block table for KV cache + page_size: Size of each page in the KV cache + + Returns: + seqlens_q_local: Query sequence lengths for local attention + cu_seqlens_q_local: Cumulative sum of query sequence lengths for local attention + seqlens_k_local: Key sequence lengths for local attention + block_table_local: Block table for local attention + """ + # Adjust attention_chunk_size based on the actual sequence length + # to avoid index out of bounds errors + max_seq_len = seq_lens_np.max() + effective_chunk_size = min(attn_chunk_size, max_seq_len) + # Make sure effective_chunk_size is divisible by page_size + effective_chunk_size = (effective_chunk_size // page_size) * page_size + if effective_chunk_size < page_size: + effective_chunk_size = page_size + attn_chunk_size = effective_chunk_size + + q_seqlens = query_start_loc_np[1:] - query_start_loc_np[:-1] + actual_batch_size = seq_lens_np.shape[0] + + # Handle if we are starting in the middle of a local attention block, + # we assume q_seqlens > 0 (for all elements), for each batch idx we compute + # the number of tokens that are not in the first local attention block and + # then we can simply use a cdiv for the rest. + # For example if we have: + # attn_chunk_size = 4 + # q_seqlens = [4, 10, 5] + # k_seqlens = [6, 17, 9] + # Then we would get: + # new_tokens_in_first_block = [2, 1, 4] + # local_blocks = [2, 4, 2] + q_tokens_in_first_block = np.minimum( + attn_chunk_size - ((seq_lens_np - q_seqlens) % attn_chunk_size), q_seqlens + ).astype(np.int32) + tokens_in_last_block = attn_chunk_size + (seq_lens_np % -attn_chunk_size) + local_blocks = 1 + cdiv(q_seqlens - q_tokens_in_first_block, attn_chunk_size) + + # Once we know the number of local blocks we can compute the request spans + # for each batch idx, we can figure out the number of "virtual" requests we + # have to make, + # For the above example we would get: + # seqlens_q_local = [2, 2, 1, 4, 4, 1, 4, 1] + # + # First Get batched arange. (E.g., [2, 4, 2] -> [0, 1, 0, 1, 2, 3, 0, 1]) + # (TODO: max a utility to share this code with _prepare_inputs) + # arange step 1. [2, 4, 2] -> [2, 6, 8] + cu_num_blocks = np.cumsum(local_blocks) + virtual_batches = cu_num_blocks[-1] + # arange step 2. [2, 6, 8] -> [0, 0, 2, 2, 2, 2, 6, 6] + block_offsets = np.repeat(cu_num_blocks - local_blocks, local_blocks) + # arange step 3. [0, 1, 0, 1, 2, 3, 0, 1] + arange = np.arange(virtual_batches, dtype=np.int32) - block_offsets + # also compute reverse arange (i.e. [1, 0, 3, 2, 1, 0, 1, 0]) + rarange = np.repeat(local_blocks, local_blocks) - arange - 1 + # Then we can compute the seqlens_q_local, handling the fact that the + # first and last blocks could be partial + seqlens_q_local = np.repeat(q_seqlens - q_tokens_in_first_block, local_blocks) + # set the first block since this may be a partial block + seqlens_q_local[arange == 0] = q_tokens_in_first_block + # set the remaining blocks + seqlens_q_local[arange > 0] = np.minimum( + seqlens_q_local - attn_chunk_size * (arange - 1), attn_chunk_size + )[arange > 0] + + # convert from q_seqlens to cu_seqlens_q + cu_seqlens_q_local = np.pad(np.cumsum(seqlens_q_local), (1, 0)).astype(np.int32) + + # compute the seqlens_k_local, + # basically a full local attention block for all but the last block in each + # batch + # For our example this will be: + # seqlens_k_local = [4, 2, 4, 4, 4, 1, 4, 1] + seqlens_k_local = np.full(cu_num_blocks[-1], attn_chunk_size, dtype=np.int32) + seqlens_k_local[cu_num_blocks - 1] = tokens_in_last_block + + k_seqstarts_absolute = np.repeat(seq_lens_np, local_blocks) - ( + rarange * attn_chunk_size + np.repeat(tokens_in_last_block, local_blocks) + ) + # For the example the local attention blocks start at: + # _b0_ _____b1_____ _b2_ + # k_seqstarts_absolute = [0, 4, 4, 8, 12, 16, 4, 8] + block_starts = k_seqstarts_absolute // page_size + + assert attn_chunk_size % page_size == 0, ( + f"attn_chunk_size {attn_chunk_size} is not " + f"divisible by page_size {page_size}" + ) + pages_per_local_batch = attn_chunk_size // page_size + + # Create a block_table for the local attention blocks + # For out example if we have a block-table like (assuming page_size=2): + # block_table = [ + # [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9], < batch 0 + # [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], < batch 1 + # [20, 21, 22, 23, 24, 25, 26, 27, 28, 29], < batch 2 + # ] + # Then for the local batches we would want a block-table like + # block_table_local = [ + # [ 0, 1 ], < local-batch 0, (batch 0, starting from k[0]) + # [ 2, 3 ], < local-batch 1, (batch 0, starting from k[4]) + # [ 12, 13 ], < local-batch 2, (batch 1, starting from k[4]) + # [ 14, 15 ], < local-batch 3, (batch 1, starting from k[8]) + # [ 16, 17 ], < local-batch 4, (batch 1, starting from k[12]) + # [ 18, 19 ], < local-batch 5, (batch 1, starting from k[16]) + # [ 22, 23 ], < local-batch 6, (batch 2, starting from k[4]) + # [ 24, 25 ], < local-batch 7, (batch 2, starting from k[8]) + # ] + block_indices = np.broadcast_to( + np.arange(pages_per_local_batch, dtype=np.int32), + (virtual_batches, pages_per_local_batch), + ) + np.expand_dims(block_starts, axis=1) + # Ensure block_indices doesn't exceed block_table dimensions + # This is a critical safety check that prevents index out of bounds errors + # when dealing with large sequences (>8192 tokens) or when the block_table + # dimensions are smaller than what would be needed for the full attention chunk size. + block_indices = block_indices.flatten().clip(max=block_table.shape[1] - 1) + batch_indices = np.repeat( + np.arange(actual_batch_size, dtype=np.int32), + local_blocks * pages_per_local_batch, + ) + block_table_local = block_table[batch_indices, block_indices].view( + virtual_batches, -1 + ) + + return seqlens_q_local, cu_seqlens_q_local, seqlens_k_local, block_table_local + + +def cdiv(a: int, b: int) -> int: + """Ceiling division.""" + return -(a // -b) + + +# TODO(hebiao064): remove this once we have a better way to handle the merge_state_v2 torch.compile issue +@torch._dynamo.disable() +def merge_state_v2_wrapper(o, s_a, o_exp, s_b): + return merge_state_v2(o, s_a, o_exp, s_b) + + +class FlashAttentionBackend(AttentionBackend): + """FlashAttention backend implementation. + + Note about the init: + - If no spec decoding + - FlashAttentionBackend will be init once when the server starts. + - If spec decoding + - FlashAttentionBackend will be init once for the target worker + - FlashAttentionMultiStepBackend will be once for the draft worker + - It will spawn num_steps FlashAttentionBackend for the draft worker + + Note about CUDA Graph: + - We only support CUDA Graph for Decode (Normal Decode and Draft Decode) and Target Verify. + - We don't support CUDA Graph for Extend and Draft Extend. + - When server init, init_cuda_graph_state will be called first and then init_cuda_graph_capture will be called. + - For each forward batch, init_replay_cuda_graph will be called first and then replay the graph. + """ + + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + speculative_step_id=0, + topk=0, + speculative_num_steps=0, + ): + super().__init__() + + assert not ( + model_runner.sliding_window_size is not None + and model_runner.model_config.is_encoder_decoder + ), "Sliding window and cross attention are not supported together" + + self.forward_metadata: FlashAttentionMetadata = None + # extra metadata for handling speculative decoding topk > 1, extended draft decode and verify + self.forward_metadata_spec_decode_expand: FlashAttentionMetadata = None + self.max_context_len = model_runner.model_config.context_len + self.device = model_runner.device + self.decode_cuda_graph_metadata = {} + self.target_verify_metadata = {} + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.kv_cache_dtype = model_runner.kv_cache_dtype + self.kv_cache_dtype_str = model_runner.server_args.kv_cache_dtype + self.page_size = model_runner.page_size + self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA + self.skip_prefill = skip_prefill + self.is_hybrid = model_runner.is_hybrid + if self.is_hybrid: + self.full_to_swa_index_mapping = ( + model_runner.token_to_kv_pool.full_to_swa_index_mapping + ) + self.topk = model_runner.server_args.speculative_eagle_topk or 0 + self.speculative_num_steps = speculative_num_steps + self.speculative_num_draft_tokens = ( + model_runner.server_args.speculative_num_draft_tokens + ) + self.speculative_step_id = speculative_step_id + + # Local attention settings + self.attention_chunk_size = ( + model_runner.attention_chunk_size + if hasattr(model_runner, "attention_chunk_size") + else None + ) + + # For each layer, the sliding_window_size can be different. This is only used for preparing SWA metadata. + # We use `layer.sliding_window_size` to decide whether to use SWA for each layer. + self.sliding_window_size = model_runner.sliding_window_size + self.has_swa = ( + self.sliding_window_size is not None and self.sliding_window_size > -1 + ) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Initialize forward metadata hence all layers in the forward pass can reuse it.""" + metadata = FlashAttentionMetadata() + seqlens_in_batch = forward_batch.seq_lens + batch_size = forward_batch.batch_size + device = seqlens_in_batch.device + + if forward_batch.forward_mode.is_decode_or_idle(): + # Draft Decode + if forward_batch.spec_info is not None: + if self.topk <= 1: + metadata.cache_seqlens_int32 = ( + seqlens_in_batch + (self.speculative_step_id + 1) + ).to(torch.int32) + metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + ( + self.speculative_step_id + 1 + ) + metadata.cu_seqlens_q = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device=device + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + else: + metadata.cache_seqlens_int32 = (seqlens_in_batch).to(torch.int32) + metadata.max_seq_len_q = self.topk + metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + metadata.cu_seqlens_q = torch.arange( + 0, + batch_size * self.topk + 1, + step=self.topk, + dtype=torch.int32, + device=device, + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + + metadata_expand = FlashAttentionMetadata() + decode_length = self.speculative_step_id + 1 + metadata_expand.cache_seqlens_int32 = torch.full( + (seqlens_in_batch.numel() * self.topk,), + decode_length, + device=device, + dtype=torch.int32, + ) + metadata_expand.max_seq_len_q = 1 + metadata_expand.cu_seqlens_q = torch.arange( + 0, + metadata_expand.cache_seqlens_int32.numel() + 1, + dtype=torch.int32, + device=device, + ) + metadata_expand.cu_seqlens_k = torch.arange( + 0, + metadata_expand.cache_seqlens_int32.numel() * decode_length + 1, + step=decode_length, + dtype=torch.int32, + device=device, + ) + # shape: [bs, num_steps, topk] -> [bs x topk, num_steps] + cache_loc = forward_batch.out_cache_loc.view( + -1, self.speculative_num_steps + ) + metadata_expand.page_table = ( + cache_loc[:, :decode_length].contiguous().to(torch.int32) + ) + self.forward_metadata_spec_decode_expand = metadata_expand + else: + # Normal Decode + metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32) + metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + metadata.cu_seqlens_q = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device=device + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0) + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + # TODO: we need to test this part for llama 4 eagle case + self._init_local_attn_metadata(forward_batch, metadata, device) + elif forward_batch.forward_mode.is_target_verify(): + if self.topk <= 1: + metadata.cache_seqlens_int32 = ( + forward_batch.seq_lens + self.speculative_num_draft_tokens + ).to(torch.int32) + metadata.max_seq_len_q = self.speculative_num_draft_tokens + metadata.max_seq_len_k = ( + forward_batch.seq_lens_cpu.max().item() + + self.speculative_num_draft_tokens + ) + metadata.cu_seqlens_q = torch.arange( + 0, + batch_size * self.speculative_num_draft_tokens + 1, + self.speculative_num_draft_tokens, + dtype=torch.int32, + device=device, + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + + self._init_local_attn_metadata(forward_batch, metadata, device) + else: + metadata.cache_seqlens_int32 = forward_batch.seq_lens.to(torch.int32) + metadata.max_seq_len_q = self.speculative_num_draft_tokens + metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + metadata.cu_seqlens_q = torch.arange( + 0, + batch_size * self.speculative_num_draft_tokens + 1, + step=self.speculative_num_draft_tokens, + dtype=torch.int32, + device=device, + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + + metadata_expand = FlashAttentionMetadata() + + metadata_expand.max_seq_len_q = 1 + metadata_expand.cu_seqlens_q = torch.arange( + 0, + forward_batch.seq_lens.numel() * self.speculative_num_draft_tokens + + 1, + dtype=torch.int32, + device=device, + ) + + # create expand page table + offsets = torch.arange( + self.speculative_num_draft_tokens, device=device + ).unsqueeze( + 0 + ) # shape: (1, self.speculative_num_draft_tokens) + cols = offsets.expand( + forward_batch.seq_lens.numel(), -1 + ) + forward_batch.seq_lens.unsqueeze(1) + cum_len = torch.nn.functional.pad( + torch.cumsum( + ( + forward_batch.seq_lens + self.speculative_num_draft_tokens + ).repeat_interleave(self.speculative_num_draft_tokens), + dim=0, + ), + (1, 0), + )[:-1] + mask_extraction_indices = ( + cols.repeat_interleave(self.speculative_num_draft_tokens, dim=0) + + cum_len[:, None] + ).view(1, -1) + mask = forward_batch.spec_info.custom_mask[ + mask_extraction_indices + ].view( + -1, self.speculative_num_draft_tokens + ) # (bsz * draft_num, draft_num) + + # shift table indices to avoid padding + # non_masked_page_table [[8, 9, 10], mask (display with int format) [[1, 0, 0], + # [8, 9, 10], [1, 1, 0], + # [8, 9, 10]] [1, 0, 1]] + # if masked with padding [[8, 0, 0], our mask without padding [[8, 9, 10], + # [8, 9, 0], [8, 9, 10], + # [8, 0, 10]] [8, 10, 9]] + # note here cache_seqlens_int32 is [1, 2, 2] so extra page indices will be ignored in each row + col_indices = offsets.expand( + mask.shape[0], self.speculative_num_draft_tokens + ) + # Build keys: if an entry is valid (mask==True), keep its original index; + # if not, add self.speculative_num_draft_tokens so that it sorts after all valid entries. + keys = torch.where( + mask, col_indices, col_indices + self.speculative_num_draft_tokens + ) + _, sort_order = torch.sort(keys, dim=1) + non_masked_page_table = ( + forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : + ] + .gather(1, cols) + .repeat_interleave(self.speculative_num_draft_tokens, dim=0) + ) # (bsz, draft_num) + metadata_expand.page_table = non_masked_page_table.gather(1, sort_order) + metadata_expand.cache_seqlens_int32 = mask.sum(dim=1).to(torch.int32) + metadata_expand.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata_expand.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) + self.forward_metadata_spec_decode_expand = metadata_expand + + if self.has_swa: + self._init_sliding_window_attn_spec_metadata( + metadata, metadata_expand + ) + + elif forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed(): + metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32) + metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0) + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + + if ( + any(forward_batch.extend_prefix_lens_cpu) + or forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND + ): + extend_seq_lens = forward_batch.extend_seq_lens + metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu) + metadata.cu_seqlens_q = torch.nn.functional.pad( + torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0) + ) + else: + metadata.max_seq_len_q = metadata.max_seq_len_k + metadata.cu_seqlens_q = metadata.cu_seqlens_k + + # Setup local attention if enabled + if forward_batch.forward_mode == ForwardMode.EXTEND: + self._init_local_attn_metadata(forward_batch, metadata, device) + + # Encoder metadata for cross attention + if forward_batch.encoder_lens is not None: + assert ( + forward_batch.encoder_lens.numel() == 1 + ), "Only encoder size 1 is supported for now" + + metadata.encoder_lens_int32 = forward_batch.encoder_lens.to(torch.int32) + metadata.encoder_cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(metadata.encoder_lens_int32, dim=0, dtype=torch.int32), + (1, 0), + ) + metadata.encoder_max_seq_len_k = metadata.encoder_lens_int32.max().item() + metadata.encoder_page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.encoder_max_seq_len_k + ] + + # Currently only support forward_batch.encoder_lens.numel() == 1 + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, + metadata.encoder_max_seq_len_k : ( + metadata.encoder_max_seq_len_k + metadata.max_seq_len_k + ), + ] + + # Convert the page table to a strided format which is needed by FA3 API + if self.page_size > 1: + self.strided_indices = torch.arange( + 0, metadata.page_table.shape[1], self.page_size, device=self.device + ) + metadata.page_table = ( + metadata.page_table[:, self.strided_indices] // self.page_size + ) + + self.forward_metadata = metadata + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + sinks: Optional[torch.Tensor] = None, + ): + if k is not None: + assert v is not None + if save_kv_cache: + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + if not self.use_mla: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, layer.k_scale, layer.v_scale + ) + else: + forward_batch.token_to_kv_pool.set_mla_kv_buffer( + layer, + cache_loc, + k, + k_rope, + ) + + # Use precomputed metadata across all layers + metadata = self.forward_metadata + + # Calculate window size (can be moved to metadata if layer properties don't change) + # we don't do layer.sliding_window_size - 1 since in model.get_attention_sliding_window_size() we already - 1 + # here is two side inclusive + is_swa = ( + layer.sliding_window_size is not None and layer.sliding_window_size > -1 + ) + window_size = (layer.sliding_window_size, 0) if is_swa else (-1, -1) + k_descale, v_descale = None, None + # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention + # has corresponding quantization method so that layer.k_scale is not None, + # 3) layer.head_dim <= 256 since fa3 kernel require fp16 and bf16 data type in this case. + if self.kv_cache_dtype_str != "auto" and layer.head_dim <= 256: + if layer.k_scale is not None: + descale_shape = (forward_batch.batch_size, layer.tp_k_head_num) + k_descale = layer.k_scale.expand(descale_shape) + v_descale = layer.v_scale.expand(descale_shape) + q = q.to(self.kv_cache_dtype) + q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None + k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None + causal = not layer.is_cross_attention + + # Check if we should use local attention + use_local_attn = ( + self.attention_chunk_size is not None + and metadata.local_attn_metadata is not None + and (hasattr(layer, "use_irope") and layer.use_irope) + ) + + # We do cascade attention for Target Verify with topk > 1 + # We don't use cascade attention for Sliding Window Attention: + # - Different window sizes should be passed in for each q in the first stage of cascade attention, but FA3 interface doesn't support pass in a list of window sizes. + # - The overhead of duplicated computation of the common prefix part is small for sliding window layers (seq_len <= window_size), so we can just expand it. + use_cascade_attn = ( + forward_batch.forward_mode.is_target_verify() + and self.topk > 1 + and not is_swa + ) + + # For fa3 interface version compatibility, we put new fields into conditional keyword args + kwargs = {} + if sinks is not None: + kwargs["sinks"] = sinks + + # Get the appropriate page table based on whether we're using local attention + if use_local_attn: + local_metadata = metadata.local_attn_metadata + page_table = local_metadata.local_block_table + cu_seqlens_q = local_metadata.local_query_start_loc + cache_seqlens = local_metadata.local_seqused_k + max_seqlen_q = local_metadata.local_max_query_len + elif is_swa and metadata.swa_spec_metadata is not None: + swa_spec_metadata = metadata.swa_spec_metadata + page_table = swa_spec_metadata.page_table + cu_seqlens_q = swa_spec_metadata.cu_seqlens_q + cache_seqlens = swa_spec_metadata.cache_seqlens_int32 + max_seqlen_q = swa_spec_metadata.max_seq_len_q + cu_seqlens_k = swa_spec_metadata.cu_seqlens_k + else: + page_table = metadata.page_table + cu_seqlens_q = metadata.cu_seqlens_q + cache_seqlens = metadata.cache_seqlens_int32 + max_seqlen_q = metadata.max_seq_len_q + cu_seqlens_k = metadata.cu_seqlens_k + + # Use Flash Attention for prefill + if not self.use_mla: + # Do multi-head attention + key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer( + layer.layer_id + ) + key_cache = key_cache.view( + -1, self.page_size, layer.tp_k_head_num, layer.head_dim + ) + value_cache = value_cache.view( + -1, self.page_size, layer.tp_v_head_num, layer.head_dim + ) + if layer.is_cross_attention: + page_table = metadata.encoder_page_table + cache_seqlens = metadata.encoder_lens_int32 + cu_seqlens_k = metadata.encoder_cu_seqlens_k + window_size = (-1, -1) + + result = flash_attn_with_kvcache( + q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache=key_cache, + v_cache=value_cache, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None, + max_seqlen_q=max_seqlen_q, + softmax_scale=layer.scaling, + causal=False if use_cascade_attn else causal, + window_size=window_size, + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=use_cascade_attn, + **kwargs, + ) + + if use_cascade_attn: + o, softmax_lse, *rest = result + o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache( + q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache=key_cache, + v_cache=value_cache, + page_table=self.forward_metadata_spec_decode_expand.page_table, + cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32, + cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q, + cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k, + max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q, + softmax_scale=layer.scaling, + causal=False, + window_size=window_size, + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=True, + **kwargs, + ) + o, _ = merge_state_v2_wrapper( + o, + softmax_lse.T.contiguous(), + o_expand, + softmax_lse_expand.T.contiguous(), + ) + else: + o = result + else: + if ( + forward_batch.attn_attend_prefix_cache is not None + and not forward_batch.forward_mode.is_target_verify() + and not forward_batch.forward_mode.is_draft_extend() + ): + # Do multi-head attention with chunked prefix cache + if forward_batch.attn_attend_prefix_cache: + assert not global_server_args_dict["disable_chunked_prefix_cache"] + # MHA for chunked prefix kv cache when running model with MLA + assert forward_batch.prefix_chunk_idx is not None + assert forward_batch.prefix_chunk_cu_seq_lens is not None + assert forward_batch.prefix_chunk_max_seq_lens is not None + + chunk_idx = forward_batch.prefix_chunk_idx + assert chunk_idx >= 0 + + assert forward_batch.mha_return_lse + output = flash_attn_varlen_func( + q=q.view(-1, layer.tp_q_head_num, layer.head_dim), + k=k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype), + v=v.view(-1, layer.tp_k_head_num, layer.v_head_dim).to(q.dtype), + cu_seqlens_q=metadata.cu_seqlens_q, + cu_seqlens_k=forward_batch.prefix_chunk_cu_seq_lens[chunk_idx], + max_seqlen_q=metadata.max_seq_len_q, + max_seqlen_k=forward_batch.prefix_chunk_max_seq_lens[chunk_idx], + softmax_scale=layer.scaling, + causal=False, + return_softmax_lse=True, + ) + else: + # MHA for extend part of sequence without attending prefix kv cache + output = flash_attn_varlen_func( + q=q.view(-1, layer.tp_q_head_num, layer.head_dim), + k=k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype), + v=v.view(-1, layer.tp_k_head_num, layer.v_head_dim).to(q.dtype), + cu_seqlens_q=metadata.cu_seqlens_q, + cu_seqlens_k=metadata.cu_seqlens_q, + max_seqlen_q=metadata.max_seq_len_q, + max_seqlen_k=metadata.max_seq_len_q, + softmax_scale=layer.scaling, + causal=True, + return_softmax_lse=forward_batch.mha_return_lse, + ) + if forward_batch.mha_return_lse: + output, lse, *rest = output + lse = torch.transpose(lse, 0, 1).contiguous() + return output, lse + return output + else: + # Do absorbed multi-latent attention + kv_cache = forward_batch.token_to_kv_pool.get_key_buffer( + layer.layer_id + ).to(q.dtype) + k_rope = kv_cache[:, :, layer.v_head_dim :] + c_kv = kv_cache[:, :, : layer.v_head_dim] + k_rope_cache = k_rope.view( + -1, + self.page_size, + layer.tp_k_head_num, + layer.head_dim - layer.v_head_dim, + ) + c_kv_cache = c_kv.view( + -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim + ) + if q_rope is not None: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + else: + q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q_nope = q_all[:, :, : layer.v_head_dim] + q_rope = q_all[:, :, layer.v_head_dim :] + + result = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=c_kv_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None, + max_seqlen_q=max_seqlen_q, + softmax_scale=layer.scaling, + causal=False if use_cascade_attn else causal, + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=use_cascade_attn, + ) + if use_cascade_attn: + o, softmax_lse, *rest = result + o_expand, softmax_lse_expand, *rest_expand = ( + flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=c_kv_cache, + qv=q_nope, + page_table=self.forward_metadata_spec_decode_expand.page_table, + cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32, + cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q, + cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k, + max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q, + softmax_scale=layer.scaling, + causal=False, + window_size=window_size, + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=True, + ) + ) + o, _ = merge_state_v2_wrapper( + o, + softmax_lse.T.contiguous(), + o_expand, + softmax_lse_expand.T.contiguous(), + ) + else: + o = result + + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + sinks: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if k is not None: + assert v is not None + if save_kv_cache: + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + if not self.use_mla: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, layer.k_scale, layer.v_scale + ) + else: + forward_batch.token_to_kv_pool.set_mla_kv_buffer( + layer, + cache_loc, + k, + k_rope, + ) + + # Use precomputed metadata across all layers + metadata = self.forward_metadata + local_attn_metadata = getattr(metadata, "local_attn_metadata", None) + use_local_attn = ( + self.attention_chunk_size is not None + and local_attn_metadata is not None + and (hasattr(layer, "use_irope") and layer.use_irope) + ) + + # When Spec Decode enabled, forward_decode would be called with two mode: + # 1. DRAFT_DECODE: we enable cascade attention when top_k > 1 + # 2. IDLE: we don’t need cascade attention, spec_info will be none in this case + use_cascade_attn = forward_batch.spec_info is not None and self.topk > 1 + + # Calculate window size (can be moved to metadata if layer properties don't change) + # we don't do layer.sliding_window_size - 1 since in model.get_attention_sliding_window_size() we already - 1 + # here is two side inclusive + window_size = ( + (layer.sliding_window_size, 0) + if layer.sliding_window_size is not None and layer.sliding_window_size > -1 + else (-1, -1) + ) + causal = not layer.is_cross_attention + + # For fa3 interface version compatibility, we put new fields into conditional keyword args + kwargs = {} + if sinks is not None: + kwargs["sinks"] = sinks + + k_descale, v_descale = None, None + # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention + # has corresponding quantization method so that layer.k_scale is not None, + # 3) layer.head_dim <= 256 since fa3 kernel require fp16 and bf16 data type in this case. + if self.kv_cache_dtype_str != "auto" and layer.head_dim <= 256: + if layer.k_scale is not None: + descale_shape = (forward_batch.batch_size, layer.tp_k_head_num) + k_descale = layer.k_scale.expand(descale_shape) + v_descale = layer.v_scale.expand(descale_shape) + q = q.to(self.kv_cache_dtype) + q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None + k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None + if not self.use_mla: + # Do multi-head attention + + key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer( + layer.layer_id + ) + key_cache = key_cache.view( + -1, self.page_size, layer.tp_k_head_num, layer.head_dim + ) + value_cache = value_cache.view( + -1, self.page_size, layer.tp_v_head_num, layer.head_dim + ) + + if layer.is_cross_attention: + # Always use non-chunked logic for cross-attention + o = flash_attn_with_kvcache( + q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache=key_cache, + v_cache=value_cache, + page_table=metadata.encoder_page_table, + cache_seqlens=metadata.encoder_lens_int32, + cu_seqlens_q=metadata.cu_seqlens_q, + cu_seqlens_k_new=metadata.encoder_cu_seqlens_k, + max_seqlen_q=1, + softmax_scale=layer.scaling, + causal=False, + window_size=(-1, -1), + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + **kwargs, + ) + elif use_local_attn: + # Use chunked (local) attention batching for self-attention + o = flash_attn_with_kvcache( + q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache=key_cache, + v_cache=value_cache, + page_table=local_attn_metadata.local_block_table, + cache_seqlens=local_attn_metadata.local_seqused_k, + cu_seqlens_q=local_attn_metadata.local_query_start_loc, + cu_seqlens_k_new=None, + max_seqlen_q=local_attn_metadata.local_max_query_len, + softmax_scale=layer.scaling, + causal=True, + window_size=(-1, -1), + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + **kwargs, + ) + else: + page_table = metadata.page_table + cache_seqlens = metadata.cache_seqlens_int32 + cu_seqlens_k = metadata.cu_seqlens_k + max_seqlen_q = metadata.max_seq_len_q + q_reshaped = q.contiguous().view( + -1, layer.tp_q_head_num, layer.head_dim + ) + + # Default: single-token self-attention + result = flash_attn_with_kvcache( + q=q_reshaped, + k_cache=key_cache, + v_cache=value_cache, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=metadata.cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + softmax_scale=layer.scaling, + causal=False if use_cascade_attn else causal, + window_size=window_size, + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=use_cascade_attn, + **kwargs, + ) + if use_cascade_attn: + o, softmax_lse, *rest = result + o_expand, softmax_lse_expand, *rest_expand = ( + flash_attn_with_kvcache( + q=q_reshaped, + k_cache=key_cache, + v_cache=value_cache, + page_table=self.forward_metadata_spec_decode_expand.page_table, + cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32, + cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q, + cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k, + max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q, + softmax_scale=layer.scaling, + causal=False, + window_size=window_size, + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=True, + **kwargs, + ) + ) + o, _ = merge_state_v2( + o, + softmax_lse.T.contiguous(), + o_expand, + softmax_lse_expand.T.contiguous(), + ) + else: + o = result + else: + # Do absorbed multi-latent attention + kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id).to( + q.dtype + ) + k_rope = kv_cache[:, :, layer.v_head_dim :] + c_kv = kv_cache[:, :, : layer.v_head_dim] + k_rope_cache = k_rope.view( + -1, + self.page_size, + layer.tp_k_head_num, + layer.head_dim - layer.v_head_dim, + ) + c_kv_cache = c_kv.view( + -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim + ) + + if q_rope is not None: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + else: + q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q_nope = q_all[:, :, : layer.v_head_dim] + q_rope = q_all[:, :, layer.v_head_dim :] + max_seqlen_q = metadata.max_seq_len_q + + result = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=c_kv_cache, + qv=q_nope, + page_table=metadata.page_table, + cache_seqlens=metadata.cache_seqlens_int32, + cu_seqlens_q=metadata.cu_seqlens_q, + cu_seqlens_k_new=metadata.cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + softmax_scale=layer.scaling, + causal=False if use_cascade_attn else causal, + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=use_cascade_attn, # softmax_lse is needed for merge states + ) + if use_cascade_attn: + o, softmax_lse, *rest = result + o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=c_kv_cache, + qv=q_nope, + page_table=self.forward_metadata_spec_decode_expand.page_table, + cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32, + cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q, + cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k, + max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q, + softmax_scale=layer.scaling, + causal=False, + window_size=window_size, + softcap=layer.logit_cap, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=True, + ) + o, _ = merge_state_v2( + o, + softmax_lse.T.contiguous(), + o_expand, + softmax_lse_expand.T.contiguous(), + ) + else: + o = result + + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + """Initialize CUDA graph state for the attention backend. + + Args: + max_bs (int): Maximum batch size to support in CUDA graphs + + This creates fixed-size tensors that will be reused during CUDA graph replay + to avoid memory allocations. + """ + max_num_pages = (self.max_context_len + self.page_size - 1) // self.page_size + + # This is being used by normal decode and draft decode when topk == 1 + self.decode_cuda_graph_metadata = { + "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device), + "cu_seqlens_q": torch.arange( + 0, max_bs + 1, dtype=torch.int32, device=self.device + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + "page_table": torch.zeros( + max_bs, + max_num_pages, + dtype=torch.int32, + device=self.device, + ), + "strided_indices": torch.arange( + 0, self.max_context_len, self.page_size, device=self.device + ), + } + # Only allocate local attention buffers if local attention is enabled + # This prevents OOM errors when local attention is not being used + if self.attention_chunk_size is not None: + # Estimate maximum sizes for local attention metadata + max_seq_len = self.max_context_len + page_size = self.page_size or 1 + attn_chunk_size = self.attention_chunk_size + max_virtual_batches = max_bs * ( + (max_seq_len + attn_chunk_size - 1) // attn_chunk_size + ) + max_pages_per_block = (attn_chunk_size + page_size - 1) // page_size + + self.decode_cuda_graph_local_attn_metadata = { + "local_query_start_loc": torch.zeros( + max_virtual_batches + 1, dtype=torch.int32, device=self.device + ), + "local_seqused_k": torch.zeros( + max_virtual_batches, dtype=torch.int32, device=self.device + ), + "local_block_table": torch.zeros( + max_virtual_batches, + max_pages_per_block, + dtype=torch.int32, + device=self.device, + ), + } + + # This is used by draft decode's first half of metadata when topk > 1 + if self.topk > 1: + self.draft_decode_metadata_topk_normal = { + "cache_seqlens": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "cu_seqlens_q": torch.arange( + 0, + max_bs * self.topk + 1, + step=self.topk, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + "page_table": torch.zeros( + max_bs, + self.max_context_len, + dtype=torch.int32, + device=self.device, + ), + } + + # This is used by draft decode's second half of metadata when topk > 1 + decode_length = self.speculative_step_id + 1 + self.draft_decode_metadata_topk_expand = { + "cache_seqlens": torch.full( + (max_bs * self.topk,), + decode_length, + device=self.device, + dtype=torch.int32, + ), + "cu_seqlens_q": torch.arange( + 0, + max_bs * self.topk + 1, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_k": torch.arange( + 0, + max_bs * self.topk * decode_length + 1, + step=decode_length, + dtype=torch.int32, + device=self.device, + ), + "page_table": torch.zeros( + max_bs * self.topk, + decode_length, + dtype=torch.int32, + device=self.device, + ), + } + + if ( + self.speculative_num_draft_tokens is not None + and self.speculative_num_draft_tokens > 0 + ): + # "page_table_draft_decode" will be set only when spec decoding enabled to save memory + self.decode_cuda_graph_metadata["page_table_draft_decode"] = torch.zeros( + max_bs, + max_num_pages, + dtype=torch.int32, + device=self.device, + ) + + self.target_verify_metadata = { + "cache_seqlens": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "cu_seqlens_q": torch.arange( + 0, + max_bs * self.speculative_num_draft_tokens + 1, + step=self.speculative_num_draft_tokens, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + "page_table": torch.zeros( + max_bs, + max_num_pages, + dtype=torch.int32, + device=self.device, + ), + "strided_indices": torch.arange( + 0, self.max_context_len, self.page_size, device=self.device + ), + } + + self.draft_extend_metadata = { + "cache_seqlens": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "cu_seqlens_q": torch.zeros( + max_bs + 1, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + "page_table": torch.zeros( + max_bs, + max_num_pages, + dtype=torch.int32, + device=self.device, + ), + "strided_indices": torch.arange( + 0, self.max_context_len, self.page_size, device=self.device + ), + } + + if self.topk > 1: + self.target_verify_metadata_topk_normal = { + "cache_seqlens": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "cu_seqlens_q": torch.arange( + 0, + max_bs * self.speculative_num_draft_tokens + 1, + step=self.speculative_num_draft_tokens, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + "page_table": torch.zeros( + max_bs, + self.max_context_len, + dtype=torch.int32, + device=self.device, + ), + } + + self.target_verify_metadata_topk_expand = { + "cache_seqlens": torch.zeros( + max_bs * self.speculative_num_draft_tokens, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_k": torch.zeros( + max_bs * self.speculative_num_draft_tokens + 1, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_q": torch.arange( + 0, + max_bs * self.speculative_num_draft_tokens + 1, + dtype=torch.int32, + device=self.device, + ), + "page_table": torch.zeros( + max_bs * self.speculative_num_draft_tokens, + self.speculative_num_draft_tokens, + dtype=torch.int32, + device=self.device, + ), + } + + if self.has_swa: + self.target_verify_metadata_topk_swa = { + "cache_seqlens": torch.zeros( + max_bs * self.speculative_num_draft_tokens, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_k": torch.zeros( + max_bs * self.speculative_num_draft_tokens + 1, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_q": torch.arange( + 0, + max_bs * self.speculative_num_draft_tokens + 1, + dtype=torch.int32, + device=self.device, + ), + "page_table": torch.zeros( + max_bs * self.speculative_num_draft_tokens, + self.max_context_len, + dtype=torch.int32, + device=self.device, + ), + } + + self.encoder_metadata = { + "encoder_page_table": torch.zeros( + max_bs, + self.max_context_len, + dtype=torch.int32, + device=self.device, + ), + "encoder_lens_int32": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "encoder_cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + } + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + """Initialize forward metadata for capturing CUDA graph.""" + metadata = FlashAttentionMetadata() + + # metadata_expand is needed for Spec Decoding when top k > 1 + metadata_expand = FlashAttentionMetadata() + + device = seq_lens.device + if forward_mode.is_decode_or_idle(): + if spec_info is not None: + # Draft Decode + if self.topk <= 1: + # When topk = 1, we use the normal decode metadata + metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[ + "cache_seqlens" + ][:bs] + metadata.max_seq_len_k = seq_lens.max().item() + ( + self.speculative_step_id + 1 + ) + metadata.cu_seqlens_q = self.decode_cuda_graph_metadata[ + "cu_seqlens_q" + ][: bs + 1] + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) + metadata.page_table = self.decode_cuda_graph_metadata[ + "page_table_draft_decode" + ][:bs, :] + self.decode_cuda_graph_metadata[bs] = metadata + else: + # When top k > 1, we need two specific draft decode metadata, and then merge states + # 1. The first half of metadata for prefix tokens + metadata.cache_seqlens_int32 = ( + self.draft_decode_metadata_topk_normal["cache_seqlens"][:bs] + ) + metadata.max_seq_len_q = self.topk + metadata.max_seq_len_k = seq_lens.max().item() + metadata.cu_seqlens_q = self.draft_decode_metadata_topk_normal[ + "cu_seqlens_q" + ][: bs + 1] + metadata.cu_seqlens_k = self.draft_decode_metadata_topk_normal[ + "cu_seqlens_k" + ][: bs + 1] + metadata.page_table = self.draft_decode_metadata_topk_normal[ + "page_table" + ][:bs, :] + + # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk) + metadata_expand.cache_seqlens_int32 = ( + self.draft_decode_metadata_topk_expand["cache_seqlens"][ + : bs * self.topk + ] + ) + metadata_expand.max_seq_len_q = 1 + metadata_expand.cu_seqlens_q = ( + self.draft_decode_metadata_topk_expand["cu_seqlens_q"][ + : bs * self.topk + 1 + ] + ) + metadata_expand.cu_seqlens_k = ( + self.draft_decode_metadata_topk_expand["cu_seqlens_k"][ + : bs * self.topk + 1 + ] + ) + metadata_expand.page_table = self.draft_decode_metadata_topk_expand[ + "page_table" + ][: bs * self.topk] + self.draft_decode_metadata_topk_normal[bs] = metadata + self.draft_decode_metadata_topk_expand[bs] = metadata_expand + else: + # Normal Decode + # Get sequence information + metadata.cache_seqlens_int32 = seq_lens.to(torch.int32) + batch_size = len(seq_lens) + device = seq_lens.device + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(seq_lens, dim=0, dtype=torch.int32), (1, 0) + ) + # Precompute maximum sequence length + metadata.max_seq_len_k = seq_lens.max().item() + # Precompute page table + metadata.page_table = self.decode_cuda_graph_metadata["page_table"][ + :bs, : + ] + # Precompute cumulative sequence lengths + metadata.cu_seqlens_q = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device=device + ) + self.decode_cuda_graph_metadata[bs] = metadata + + if self.attention_chunk_size is not None: + self._update_local_attn_metadata_for_capture(metadata, batch_size) + + elif forward_mode.is_target_verify(): + if self.topk <= 1: + metadata.cache_seqlens_int32 = self.target_verify_metadata[ + "cache_seqlens" + ][:bs] + metadata.cache_seqlens_int32.copy_( + (seq_lens + self.speculative_num_draft_tokens) + ) + + metadata.max_seq_len_q = self.speculative_num_draft_tokens + metadata.max_seq_len_k = ( + seq_lens.max().item() + self.speculative_num_draft_tokens + ) + + metadata.cu_seqlens_q = torch.arange( + 0, + bs * self.speculative_num_draft_tokens + 1, + self.speculative_num_draft_tokens, + dtype=torch.int32, + device=device, + ) + + metadata.cu_seqlens_k = self.target_verify_metadata["cu_seqlens_k"][ + : (bs + 1) + ] + + metadata.page_table = self.target_verify_metadata["page_table"][:bs, :] + + self.target_verify_metadata[bs] = metadata + else: + # When topk > 1, we need two specific target verify metadata, and then merge states + # 1. The first half of metadata for prefix tokens + metadata.cache_seqlens_int32 = self.target_verify_metadata_topk_normal[ + "cache_seqlens" + ][:bs] + metadata.max_seq_len_q = self.speculative_num_draft_tokens + # metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item(), do this in replay + metadata.cu_seqlens_q = self.target_verify_metadata_topk_normal[ + "cu_seqlens_q" + ][: bs + 1] + metadata.cu_seqlens_k = self.target_verify_metadata_topk_normal[ + "cu_seqlens_k" + ][: bs + 1] + metadata.page_table = self.target_verify_metadata_topk_normal[ + "page_table" + ][:bs, :] + + # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk) + metadata_expand.cache_seqlens_int32 = ( + self.target_verify_metadata_topk_expand["cache_seqlens"][ + : bs * self.speculative_num_draft_tokens + ] + ) + metadata_expand.max_seq_len_q = 1 + metadata_expand.cu_seqlens_q = self.target_verify_metadata_topk_expand[ + "cu_seqlens_q" + ][: bs * self.speculative_num_draft_tokens + 1] + metadata_expand.cu_seqlens_k = self.target_verify_metadata_topk_expand[ + "cu_seqlens_k" + ][: bs * self.speculative_num_draft_tokens + 1] + + metadata_expand.page_table = self.target_verify_metadata_topk_expand[ + "page_table" + ][: bs * self.speculative_num_draft_tokens] + + self.target_verify_metadata_topk_normal[bs] = metadata + self.target_verify_metadata_topk_expand[bs] = metadata_expand + + if self.has_swa: + metadata_swa = FlashAttentionMetadata() + metadata_swa.cache_seqlens_int32 = ( + self.target_verify_metadata_topk_swa["cache_seqlens"][ + : bs * self.speculative_num_draft_tokens + ] + ) + metadata_swa.max_seq_len_q = 1 + metadata_swa.cu_seqlens_q = self.target_verify_metadata_topk_swa[ + "cu_seqlens_q" + ][: bs * self.speculative_num_draft_tokens + 1] + metadata_swa.cu_seqlens_k = self.target_verify_metadata_topk_swa[ + "cu_seqlens_k" + ][: bs * self.speculative_num_draft_tokens + 1] + + metadata_swa.page_table = self.target_verify_metadata_topk_swa[ + "page_table" + ][: bs * self.speculative_num_draft_tokens] + self.target_verify_metadata_topk_swa[bs] = metadata_swa + metadata.swa_spec_metadata = metadata_swa + + elif forward_mode.is_draft_extend(): + metadata.cache_seqlens_int32 = self.draft_extend_metadata["cache_seqlens"][ + :bs + ] + metadata.cache_seqlens_int32.copy_(seq_lens) + + num_tokens_per_bs = num_tokens // bs + metadata.max_seq_len_q = num_tokens_per_bs + metadata.max_seq_len_k = seq_lens.max().item() + + metadata.cu_seqlens_q = torch.arange( + 0, + bs * num_tokens_per_bs + 1, + num_tokens_per_bs, + dtype=torch.int32, + device=device, + ) + + metadata.cu_seqlens_k = self.draft_extend_metadata["cu_seqlens_k"][ + : (bs + 1) + ] + metadata.page_table = self.draft_extend_metadata["page_table"][:bs, :] + + self.draft_extend_metadata[bs] = metadata + + if encoder_lens is not None: + encoder_bs = encoder_lens.numel() + metadata.encoder_lens_int32 = self.encoder_metadata["encoder_lens_int32"][ + :encoder_bs + ] + metadata.encoder_cu_seqlens_k = self.encoder_metadata[ + "encoder_cu_seqlens_k" + ][: (encoder_bs + 1)] + + metadata.encoder_page_table = self.encoder_metadata["encoder_page_table"][ + :bs, : + ] + + self.forward_metadata = metadata + self.forward_metadata_spec_decode_expand = metadata_expand + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + out_cache_loc: Optional[torch.Tensor] = None, + ): + """Initialize forward metadata for replaying CUDA graph.""" + seq_lens = seq_lens[:bs] + seq_lens_cpu = seq_lens_cpu[:bs] + req_pool_indices = req_pool_indices[:bs] + device = seq_lens.device + metadata = None + metadata_expand = None + + if forward_mode.is_decode_or_idle(): + + if spec_info is not None: + # Draft Decode + if self.topk <= 1: + # When topk = 1, we use the normal decode metadata + metadata = self.decode_cuda_graph_metadata[bs] + max_len = seq_lens_cpu.max().item() + metadata.max_seq_len_k = max_len + self.speculative_step_id + 1 + max_seq_pages = ( + metadata.max_seq_len_k + self.page_size - 1 + ) // self.page_size + + normal_decode_set_metadata( + metadata.cache_seqlens_int32, + metadata.cu_seqlens_k, + metadata.page_table, + self.req_to_token, + req_pool_indices, + self.decode_cuda_graph_metadata["strided_indices"], + max_seq_pages, + seq_lens, + self.speculative_step_id + 1, + self.page_size, + ) + + else: + # When top k > 1, we need two specific draft decode metadata, and then merge states + # 1. The first half of metadata for prefix tokens + metadata = self.draft_decode_metadata_topk_normal[bs] + metadata.cache_seqlens_int32.copy_(seq_lens) + # metadata.max_seq_len_q = self.topk, already set in capture + metadata.max_seq_len_k = seq_lens_cpu.max().item() + # metadata.cu_seqlens_q already set in capture + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ) + ) + + page_table = self.req_to_token[ + req_pool_indices, : metadata.max_seq_len_k + ] + + metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table) + + # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk) + metadata_expand = self.draft_decode_metadata_topk_expand[bs] + decode_length = self.speculative_step_id + 1 + # shape: [bs, num_steps, topk] -> [bs x topk, num_steps] + cache_loc = out_cache_loc.view(-1, self.speculative_num_steps) + metadata_expand.page_table[: cache_loc.shape[0]].copy_( + cache_loc[:, :decode_length] + ) + # TODO: Handle local attention metadata for draft decode when llama4 eagle is supported + else: + # Normal Decode + metadata = self.decode_cuda_graph_metadata[bs] + max_len = seq_lens_cpu.max().item() + max_seq_pages = (max_len + self.page_size - 1) // self.page_size + metadata.max_seq_len_k = max_len + + normal_decode_set_metadata( + metadata.cache_seqlens_int32, + metadata.cu_seqlens_k, + metadata.page_table, + self.req_to_token, + req_pool_indices, + self.decode_cuda_graph_metadata["strided_indices"], + max_seq_pages, + seq_lens, + 0, + self.page_size, + ) + + self._update_local_attn_metadata_for_replay( + metadata, + bs, + ) + elif forward_mode.is_target_verify(): + if self.topk <= 1: + metadata = self.target_verify_metadata[bs] + metadata.cache_seqlens_int32.copy_( + (seq_lens + self.speculative_num_draft_tokens) + ) + + metadata.max_seq_len_k = ( + seq_lens_cpu.max().item() + self.speculative_num_draft_tokens + ) + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32) + ) + max_seq_pages = ( + metadata.max_seq_len_k + self.page_size - 1 + ) // self.page_size + page_indices = self.req_to_token[ + req_pool_indices[:, None], + self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages], + ] + page_indices //= self.page_size + metadata.page_table[:, :max_seq_pages].copy_(page_indices) + else: + # When topk > 1, we need two specific target verify metadata, and then merge states + # 1. The first half of metadata for prefix tokens + metadata = self.target_verify_metadata_topk_normal[bs] + metadata.cache_seqlens_int32.copy_(seq_lens) + # metadata.max_seq_len_q = self.speculative_num_draft_tokens, already set in capture + metadata.max_seq_len_k = seq_lens_cpu.max().item() + # metadata.cu_seqlens_q already set in capture + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32) + ) + page_table = self.req_to_token[ + req_pool_indices, : metadata.max_seq_len_k + ] + metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table) + + # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk) + metadata_expand = self.target_verify_metadata_topk_expand[bs] + + # metadata_expand.max_seq_len_q = 1, already set in capture + # metadata_expand.cu_seqlens_q already set in capture + offsets = torch.arange( + self.speculative_num_draft_tokens, device=device + ).unsqueeze( + 0 + ) # shape: (1, self.speculative_num_draft_tokens) + + cols = offsets.expand(seq_lens.numel(), -1) + seq_lens.unsqueeze(1) + cum_len = torch.nn.functional.pad( + torch.cumsum( + ( + seq_lens + self.speculative_num_draft_tokens + ).repeat_interleave(self.speculative_num_draft_tokens), + dim=0, + ), + (1, 0), + )[:-1] + mask_extraction_indices = ( + cols.repeat_interleave(self.speculative_num_draft_tokens, dim=0) + + cum_len[:, None] + ).view(1, -1) + # avoid extracting padded seq indices which will be out of boundary + mask_extraction_indices[ + :, + spec_info.positions.numel() * self.speculative_num_draft_tokens :, + ].fill_(0) + mask = spec_info.custom_mask[mask_extraction_indices].view( + -1, self.speculative_num_draft_tokens + ) # (bsz * draft_num, draft_num) + + col_indices = offsets.expand( + mask.shape[0], self.speculative_num_draft_tokens + ) + keys = torch.where( + mask, + col_indices, + col_indices + self.speculative_num_draft_tokens, + ) + _, sort_order = torch.sort(keys, dim=1) + + non_masked_page_table = ( + self.req_to_token[req_pool_indices, :] + .gather(1, cols) + .repeat_interleave(self.speculative_num_draft_tokens, dim=0) + ) # (bsz, draft_num) + + metadata_expand.page_table.copy_( + non_masked_page_table.gather(1, sort_order) + ) + metadata_expand.cache_seqlens_int32.copy_(mask.sum(dim=1)) + metadata_expand.cu_seqlens_k[1:].copy_( + torch.cumsum( + metadata_expand.cache_seqlens_int32, + dim=0, + dtype=torch.int32, + ) + ) + + if self.has_swa: + metadata_swa = self.target_verify_metadata_topk_swa[bs] + self._init_sliding_window_attn_spec_metadata( + metadata, metadata_expand, metadata_swa + ) + + elif forward_mode.is_draft_extend(): + metadata = self.draft_extend_metadata[bs] + metadata.cache_seqlens_int32.copy_(seq_lens) + + metadata.max_seq_len_k = seq_lens_cpu.max().item() + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32) + ) + accept_length = spec_info.accept_length[:bs] + if spec_info.accept_length_cpu: + metadata.max_seq_len_q = max(spec_info.accept_length_cpu) + 1 + else: + metadata.max_seq_len_q = 1 + + metadata.cu_seqlens_q[1:].copy_( + torch.cumsum(accept_length, dim=0, dtype=torch.int32) + ) + + max_seq_pages = ( + metadata.max_seq_len_k + self.page_size - 1 + ) // self.page_size + page_indices = self.req_to_token[ + req_pool_indices[:, None], + self.draft_extend_metadata["strided_indices"][:max_seq_pages], + ] + metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size) + + if encoder_lens is not None: + # Only support encoder size 1 for now + metadata.encoder_max_seq_len_k = encoder_lens[0] + metadata.encoder_lens_int32.copy_(encoder_lens[:1]) + metadata.encoder_cu_seqlens_k[1:].copy_( + torch.cumsum(metadata.encoder_lens_int32, dim=0, dtype=torch.int32) + ) + + metadata.encoder_page_table[:, : metadata.encoder_max_seq_len_k].copy_( + self.req_to_token[req_pool_indices, : metadata.encoder_max_seq_len_k] + ) + + # Update the regular page table + page_table = self.req_to_token[ + req_pool_indices, + metadata.encoder_max_seq_len_k : ( + metadata.encoder_max_seq_len_k + metadata.max_seq_len_k + ), + ] + metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table) + + self.forward_metadata = metadata + self.forward_metadata_spec_decode_expand = metadata_expand + + def get_cuda_graph_seq_len_fill_value(self): + """Get the fill value for sequence length in CUDA graph.""" + return 1 + + def _init_local_attn_metadata( + self, forwardbatch: ForwardBatch, metadata: FlashAttentionMetadata, device + ): + """Centralized utility to initialize local_attn_metadata if chunked attention is enabled.""" + if self.attention_chunk_size is None: + metadata.local_attn_metadata = None + return + + cu_seqlens_q = metadata.cu_seqlens_q + cache_seqlens_int32 = metadata.cache_seqlens_int32 + if self.is_hybrid: + page_table = self.full_to_swa_index_mapping[metadata.page_table].to( + torch.int32 + ) + else: + page_table = metadata.page_table + if cu_seqlens_q is None or cache_seqlens_int32 is None or page_table is None: + metadata.local_attn_metadata = None + return + + cu_seqlens_q_np = cu_seqlens_q.cpu().numpy() + seq_lens_np = cache_seqlens_int32.cpu().numpy() + ( + seqlens_q_local_np, + cu_seqlens_q_local_np, + seqlens_k_local_np, + block_table_local, + ) = make_local_attention_virtual_batches( + self.attention_chunk_size, + cu_seqlens_q_np, + seq_lens_np, + page_table, + self.page_size, + ) + + local_metadata = FlashAttentionMetadata.LocalAttentionMetadata( + local_query_start_loc=torch.from_numpy(cu_seqlens_q_local_np).to(device), + local_seqused_k=torch.from_numpy(seqlens_k_local_np).to(device), + local_block_table=block_table_local.to(device), + local_max_query_len=int(seqlens_q_local_np.max()), + local_max_seq_len=int(seqlens_k_local_np.max()), + ) + metadata.local_attn_metadata = local_metadata + + def _update_local_attn_metadata_for_capture( + self, metadata: FlashAttentionMetadata, bs: int + ): + """Update local attention metadata during CUDA graph capture phase. + + This method calculates the exact buffer sizes needed for local attention metadata + during the CUDA graph capture phase, optimizing memory usage by creating views of + pre-allocated buffers with exactly the sizes needed. + """ + seq_lens_capture = metadata.cache_seqlens_int32 + max_seq_len = int(seq_lens_capture.max().item()) + page_table_capture = metadata.page_table + + cu_seqlens_q_np = metadata.cu_seqlens_q.cpu().numpy() + seqlens_np = seq_lens_capture.cpu().numpy() + ( + seqlens_q_local_np, + cu_seqlens_q_local_np, + seqlens_k_local_np, + block_table_local_np, + ) = make_local_attention_virtual_batches( + self.attention_chunk_size, + cu_seqlens_q_np, + seqlens_np, + page_table_capture, + self.page_size, + ) + + # Get exact dimensions from the calculation + q_len = len(cu_seqlens_q_local_np) + k_len = len(seqlens_k_local_np) + b0 = block_table_local_np.shape[0] if block_table_local_np.shape[0] > 0 else bs + b1 = block_table_local_np.shape[1] if block_table_local_np.shape[1] > 0 else 1 + + # Create views of the pre-allocated buffers with exactly these sizes + # This is the key optimization - we only use the memory we actually need + local_query_start_loc = self.decode_cuda_graph_local_attn_metadata[ + "local_query_start_loc" + ][:q_len] + + local_seqused_k = self.decode_cuda_graph_local_attn_metadata["local_seqused_k"][ + :k_len + ] + + local_block_table = self.decode_cuda_graph_local_attn_metadata[ + "local_block_table" + ][:b0, :b1] + + metadata.local_attn_metadata = FlashAttentionMetadata.LocalAttentionMetadata( + local_query_start_loc=local_query_start_loc, + local_seqused_k=local_seqused_k, + local_block_table=local_block_table, + local_max_query_len=1, + local_max_seq_len=max_seq_len, + ) + + def _update_local_attn_metadata_for_replay( + self, + metadata: FlashAttentionMetadata, + bs: int, + ): + """Update preallocated local attention metadata in-place before CUDA graph replay.""" + if self.attention_chunk_size is None: + return + + # Access preallocated buffers + local_q_buf = self.decode_cuda_graph_local_attn_metadata[ + "local_query_start_loc" + ] + local_k_buf = self.decode_cuda_graph_local_attn_metadata["local_seqused_k"] + local_block_buf = self.decode_cuda_graph_local_attn_metadata[ + "local_block_table" + ] + cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"] + + # Create a modified version for local attention that only processes the last token + # This mimics the normal decode pattern + cu_seqlens_q = torch.arange( + bs + 1, device=cu_seqlens_q.device, dtype=cu_seqlens_q.dtype + ) + seqlens = metadata.cache_seqlens_int32[:bs] + # Slice the page_table to match the batch size and actual sequence length + # This serves three important purposes: + # 1. Ensures we only process the actual batch size (bs) and not the maximum batch size + # 2. Limits the sequence length to prevent processing padding tokens or garbage values + # 3. Prevents zeros in the block table which can cause garbage output during replay + # + # Without this slicing, the pre-allocated page_table may contain zeros or invalid indices + # beyond the actual sequence length, leading to incorrect attention calculations + max_seq_len = int(seqlens.max().item()) + if self.is_hybrid: + sliced_page_table = self.full_to_swa_index_mapping[ + metadata.page_table[:bs, :max_seq_len] + ].to(torch.int32) + else: + sliced_page_table = metadata.page_table[:bs, :max_seq_len] + + cu_seqlens_q_np = cu_seqlens_q.cpu().numpy() + seqlens_np = seqlens.cpu().numpy() + ( + seqlens_q_local_np, + cu_seqlens_q_local_np, + seqlens_k_local_np, + block_table_local, + ) = make_local_attention_virtual_batches( + self.attention_chunk_size, + cu_seqlens_q_np, + seqlens_np, + sliced_page_table, + self.page_size, + ) + + # Convert back to tensors + device = local_q_buf.device + cu_seqlens_q_local = torch.from_numpy(cu_seqlens_q_local_np).to(device) + seqlens_k_local = torch.from_numpy(seqlens_k_local_np).to(device) + block_table_local = block_table_local.to(device) + # Get sizes + q_len = cu_seqlens_q_local.shape[0] + k_len = seqlens_k_local.shape[0] + b0, b1 = block_table_local.shape + + # In-place updates into preallocated tensors and zero out the unused space + local_q_buf[:q_len].copy_(cu_seqlens_q_local) + local_q_buf[q_len:].fill_(0) + local_k_buf[:k_len].copy_(seqlens_k_local) + local_k_buf[k_len:].fill_(0) + local_block_buf[:b0, :b1].copy_(block_table_local) + local_block_buf[b0:, :].fill_(0) + local_block_buf[:b0, b1:].fill_(0) + + if metadata.local_attn_metadata is not None: + lam = metadata.local_attn_metadata + lam.local_max_query_len = int(seqlens_q_local_np.max()) + lam.local_max_seq_len = int(seqlens_k_local_np.max()) + + def _init_sliding_window_attn_spec_metadata( + self, + metadata: FlashAttentionMetadata, + metadata_expand: FlashAttentionMetadata, + metadata_swa: Optional[FlashAttentionMetadata] = None, + ): + # TODO: support page_size > 1 for swa spec + assert ( + self.page_size == 1 + ), "FlashAttention backend doesn't support topk > 1 speculative decoding with page size > 1 sliding window attention" + + cache_seqlens_int32 = ( + metadata.cache_seqlens_int32.repeat_interleave( + self.speculative_num_draft_tokens + ) + + metadata_expand.cache_seqlens_int32 + ) + cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32), (1, 0) + ) + bs = cache_seqlens_int32.shape[0] + page_table = ( + metadata.page_table.new_zeros( + (bs, metadata.max_seq_len_k + metadata_expand.page_table.shape[1]) + ) + if metadata_swa is None + else metadata_swa.page_table + ) + + prepare_swa_spec_page_table_triton( + page_table, + metadata.page_table, + metadata_expand.page_table, + metadata.cache_seqlens_int32, + metadata_expand.cache_seqlens_int32, + self.speculative_num_draft_tokens, + ) + + if metadata_swa is None: + metadata_swa = FlashAttentionMetadata() + metadata_swa.max_seq_len_q = 1 + metadata_swa.cu_seqlens_q = metadata_expand.cu_seqlens_q + metadata_swa.cache_seqlens_int32 = cache_seqlens_int32 + metadata_swa.cu_seqlens_k = cu_seqlens_k + metadata_swa.page_table = page_table + else: + metadata_swa.cache_seqlens_int32.copy_(cache_seqlens_int32) + metadata_swa.cu_seqlens_k.copy_(cu_seqlens_k) + + metadata.swa_spec_metadata = metadata_swa + + +@triton.jit +def _prepare_swa_spec_page_table_kernel( + dst_ptr, + src_a_ptr, + src_b_ptr, + seq_len_a_ptr, + seq_len_b_ptr, + dst_stride_m, + dst_stride_n, + a_stride_m, + a_stride_n, + b_stride_m, + b_stride_n, + LEN_A: tl.constexpr, + LEN_B: tl.constexpr, + REPEAT_STEP: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + idx_a = pid_m // REPEAT_STEP + idx_b = pid_m + seq_len_a = tl.load(seq_len_a_ptr + idx_a) + seq_len_b = tl.load(seq_len_b_ptr + idx_b) + + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + total_len = seq_len_a + seq_len_b + + if pid_n * BLOCK_N >= total_len: + return + + mask = offs_n < total_len + dst = dst_ptr + pid_m * dst_stride_m + offs_n * dst_stride_n + + if (pid_n + 1) * BLOCK_N < seq_len_a: + a_ptr = src_a_ptr + idx_a * a_stride_m + offs_n * a_stride_n + a_mask = mask & (offs_n < LEN_A) + val = tl.load(a_ptr, mask=a_mask, other=0) + tl.store(dst, val, mask=mask) + elif pid_n * BLOCK_N >= seq_len_a: + offs_b = offs_n - seq_len_a + b_ptr = src_b_ptr + idx_b * b_stride_m + offs_b * b_stride_n + b_mask = mask & (offs_b < LEN_B) + val = tl.load(b_ptr, mask=b_mask, other=0) + tl.store(dst, val, mask=mask) + else: + # mixed part + a_offs = offs_n + a_mask = (a_offs < seq_len_a) & (a_offs < LEN_A) + a_ptr = src_a_ptr + idx_a * a_stride_m + a_offs * a_stride_n + a_val = tl.load(a_ptr, mask=a_mask, other=0) + + b_offs = offs_n - seq_len_a + b_mask = (b_offs >= 0) & (b_offs < seq_len_b) & (b_offs < LEN_B) + b_ptr = src_b_ptr + idx_b * b_stride_m + b_offs * b_stride_n + b_val = tl.load(b_ptr, mask=b_mask, other=0) + + result = tl.where(offs_n < seq_len_a, a_val, b_val) + tl.store(dst, result, mask=mask) + + +def prepare_swa_spec_page_table_triton( + page_table_dst: torch.Tensor, + page_table_a: torch.Tensor, + page_table_b: torch.Tensor, # expand page table + seq_len_a: torch.Tensor, + seq_len_b: torch.Tensor, # expand seq lens + speculative_num_draft_tokens: int, +): + # concat page_table and expand page_table by kv seq length + bs = seq_len_a.numel() + bs_expand = seq_len_b.numel() + assert bs_expand == bs * speculative_num_draft_tokens + + LEN_A = page_table_a.shape[1] + LEN_B = page_table_b.shape[1] + LEN_OUT = LEN_A + LEN_B + REPEAT_STEP = speculative_num_draft_tokens + BLOCK_N = 256 + + grid = (bs_expand, triton.cdiv(LEN_OUT, BLOCK_N)) + _prepare_swa_spec_page_table_kernel[grid]( + page_table_dst, + page_table_a, + page_table_b, + seq_len_a, + seq_len_b, + page_table_dst.stride(0), + page_table_dst.stride(1), + page_table_a.stride(0), + page_table_a.stride(1), + page_table_b.stride(0), + page_table_b.stride(1), + LEN_A=LEN_A, + LEN_B=LEN_B, + REPEAT_STEP=REPEAT_STEP, + BLOCK_N=BLOCK_N, + num_warps=4, + ) + + +class FlashAttentionMultiStepBackend: + + def __init__( + self, model_runner: ModelRunner, topk: int, speculative_num_steps: int + ): + self.model_runner = model_runner + self.topk = topk + self.speculative_num_steps = speculative_num_steps + self.attn_backends = [] + for i in range(self.speculative_num_steps): + self.attn_backends.append( + FlashAttentionBackend( + model_runner, + speculative_step_id=i, + topk=self.topk, + speculative_num_steps=self.speculative_num_steps, + ) + ) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata(forward_batch) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + for i in range(self.speculative_num_steps): + self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens) + + def init_forward_metadata_capture_cuda_graph( + self, + forward_batch: ForwardBatch, + ): + assert forward_batch.spec_info is not None + assert isinstance(forward_batch.spec_info, EagleDraftInput) + + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata_capture_cuda_graph( + forward_batch.batch_size, + forward_batch.batch_size * self.topk, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + encoder_lens=forward_batch.encoder_lens, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + ) + + def init_forward_metadata_replay_cuda_graph( + self, forward_batch: ForwardBatch, bs: int + ): + assert forward_batch.spec_info is not None + assert isinstance(forward_batch.spec_info, EagleDraftInput) + + for i in range(self.speculative_num_steps - 1): + # TODO: incrementally update the metadata for the later steps, + # so that they do not need to recompute everything from scratch. + self.attn_backends[i].init_forward_metadata_replay_cuda_graph( + bs, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + encoder_lens=forward_batch.encoder_lens, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + seq_lens_cpu=forward_batch.seq_lens_cpu, + out_cache_loc=forward_batch.out_cache_loc, + ) + + +# @torch.compile(dynamic=True, backend=get_compiler_backend()) +# TODO: fuse these kernels +# NOTE: torch.compile makes it slower in speculative decoding +def normal_decode_set_metadata( + cache_seqlens_int32: torch.Tensor, + cu_seqlens_k: torch.Tensor, + page_table: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + strided_indices: torch.Tensor, + max_seq_pages: torch.Tensor, + seq_lens: torch.Tensor, + seq_len_delta: int, + page_size: int, +): + cache_seqlens_int32.copy_(seq_lens + seq_len_delta) + cu_seqlens_k[1:].copy_(torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32)) + page_indices = req_to_token[ + req_pool_indices[:, None], + strided_indices[:max_seq_pages][None, :], + ] + page_table[:, :max_seq_pages].copy_(page_indices // page_size) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py new file mode 100644 index 000000000..a5b207c77 --- /dev/null +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -0,0 +1,1441 @@ +from __future__ import annotations + +""" +Support different attention backends. +Now there are two backends: FlashInfer and Triton. +FlashInfer is faster and Triton is easier to customize. +Each backend supports two operators: extend (i.e. prefill with cached prefix) and decode. +""" + +import os +from dataclasses import dataclass +from enum import Enum, auto +from functools import partial +from typing import TYPE_CHECKING, Callable, List, Optional, Union + +import torch + +if os.environ["SGLANG_ENABLE_TORCH_COMPILE"] == "1": + import logging + + torch._logging.set_logs(dynamo=logging.ERROR) + torch._dynamo.config.suppress_errors = True + +from sglang.global_config import global_config +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.layers.radix_attention import AttentionType +from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput +from sglang.srt.utils import ( + is_flashinfer_available, + is_sm100_supported, + next_power_of_2, +) + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + +if is_flashinfer_available(): + from flashinfer import ( + BatchDecodeWithPagedKVCacheWrapper, + BatchPrefillWithPagedKVCacheWrapper, + BatchPrefillWithRaggedKVCacheWrapper, + ) + from flashinfer.cascade import merge_state + from flashinfer.decode import _get_range_buf, get_seq_lens + + +class WrapperDispatch(Enum): + SLIDING_WINDOW = auto() + CROSS_ATTENTION = auto() + + +@dataclass +class DecodeMetadata: + decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper] + + +@dataclass +class PrefillMetadata: + prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper] + use_ragged: bool + extend_no_prefix: bool + + +# Reuse this workspace buffer across all flashinfer wrappers +global_workspace_buffer = None + +# Use as a fast path to override the indptr in flashinfer's plan function +# This is used to remove some host-to-device copy overhead. +global_override_indptr_cpu = None + + +class FlashInferAttnBackend(AttentionBackend): + """Flashinfer attention kernels.""" + + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + kv_indptr_buf: Optional[torch.Tensor] = None, + kv_last_page_len_buf: Optional[torch.Tensor] = None, + ): + super().__init__() + + # Parse constants + self.decode_use_tensor_cores = should_use_tensor_core( + kv_cache_dtype=model_runner.kv_cache_dtype, + num_attention_heads=model_runner.model_config.num_attention_heads + // get_attention_tp_size(), + num_kv_heads=model_runner.model_config.get_num_kv_heads( + get_attention_tp_size() + ), + ) + self.max_context_len = model_runner.model_config.context_len + self.skip_prefill = skip_prefill + self.is_multimodal = model_runner.model_config.is_multimodal + + assert not ( + model_runner.sliding_window_size is not None + and model_runner.model_config.is_encoder_decoder + ), "Sliding window and cross attention are not supported together" + + if model_runner.sliding_window_size is not None: + self.num_wrappers = 2 + self.dispatch_reason = WrapperDispatch.SLIDING_WINDOW + elif model_runner.model_config.is_encoder_decoder: + self.num_wrappers = 2 + self.dispatch_reason = WrapperDispatch.CROSS_ATTENTION + else: + self.num_wrappers = 1 + self.dispatch_reason = None + + # Qwen2/Qwen3 models require higher flashinfer workspace size + if ( + "Qwen2ForCausalLM" in model_runner.model_config.hf_config.architectures + or "Qwen3ForCausalLM" in model_runner.model_config.hf_config.architectures + or "MiMoForCausalLM" in model_runner.model_config.hf_config.architectures + ): + global_config.flashinfer_workspace_size = 512 * 1024 * 1024 + + # Allocate buffers + global global_workspace_buffer + if global_workspace_buffer is None: + # different from flashinfer zero_init_global_workspace_buffer + global_workspace_buffer = torch.empty( + global_config.flashinfer_workspace_size, + dtype=torch.uint8, + device=model_runner.device, + ) + self.workspace_buffer = global_workspace_buffer + max_bs = model_runner.req_to_token_pool.size + if kv_indptr_buf is None: + self.kv_indptr = [ + torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + for _ in range(self.num_wrappers) + ] + else: + assert self.num_wrappers == 1 + self.kv_indptr = [kv_indptr_buf] + + if kv_last_page_len_buf is None: + self.kv_last_page_len = torch.ones( + (max_bs,), dtype=torch.int32, device=model_runner.device + ) + else: + assert self.num_wrappers == 1 + self.kv_last_page_len = kv_last_page_len_buf + + if not self.skip_prefill: + self.qo_indptr = [ + torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + for _ in range(self.num_wrappers) + ] + + fmha_backend = "auto" + if is_sm100_supported(): + fmha_backend = "cutlass" + self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper( + self.workspace_buffer, "NHD", backend=fmha_backend + ) + + # Two wrappers: one for sliding window attention and one for full attention. + # Using two wrappers is unnecessary in the current PR, but are prepared for future PRs + self.prefill_wrappers_paged = [] + self.prefill_wrappers_verify = [] + self.decode_wrappers = [] + for _ in range(self.num_wrappers): + if not skip_prefill: + self.prefill_wrappers_paged.append( + BatchPrefillWithPagedKVCacheWrapper( + self.workspace_buffer, + "NHD", + backend="fa2", + ) + ) + self.prefill_wrappers_verify.append( + BatchPrefillWithPagedKVCacheWrapper( + self.workspace_buffer, + "NHD", + ) + ) + self.decode_wrappers.append( + BatchDecodeWithPagedKVCacheWrapper( + self.workspace_buffer, + "NHD", + use_tensor_cores=self.decode_use_tensor_cores, + ) + ) + + # Create indices updater + if not skip_prefill: + self.indices_updater_prefill = FlashInferIndicesUpdaterPrefill( + model_runner, self + ) # for verify + self.indices_updater_decode = FlashInferIndicesUpdaterDecode(model_runner, self) + + # Other metadata + self.forward_metadata: Union[PrefillMetadata, DecodeMetadata] = None + self.decode_cuda_graph_metadata = {} + self.prefill_cuda_graph_metadata = {} # For verify + self.draft_extend_cuda_graph_metadata = {} # For draft extend + + def init_forward_metadata(self, forward_batch: ForwardBatch): + if forward_batch.forward_mode.is_decode_or_idle(): + self.indices_updater_decode.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_cpu, + forward_batch.seq_lens_sum, + decode_wrappers=self.decode_wrappers, + encoder_lens=forward_batch.encoder_lens, + spec_info=forward_batch.spec_info, + ) + self.forward_metadata = DecodeMetadata(self.decode_wrappers) + elif forward_batch.forward_mode.is_draft_extend(): + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_cpu, + forward_batch.seq_lens_sum, + prefix_lens=None, + prefill_wrappers=self.prefill_wrappers_paged, + use_ragged=False, + encoder_lens=forward_batch.encoder_lens, + spec_info=forward_batch.spec_info, + ) + self.forward_metadata = PrefillMetadata( + self.prefill_wrappers_paged, False, False + ) + elif forward_batch.forward_mode.is_target_verify(): + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_cpu, + forward_batch.seq_lens_sum, + prefix_lens=None, + prefill_wrappers=self.prefill_wrappers_verify, + use_ragged=False, + encoder_lens=forward_batch.encoder_lens, + spec_info=forward_batch.spec_info, + ) + self.forward_metadata = PrefillMetadata( + self.prefill_wrappers_verify, False, False + ) + else: + prefix_lens = forward_batch.extend_prefix_lens + + if self.is_multimodal: + use_ragged = False + extend_no_prefix = False + else: + use_ragged = True + extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu) + + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_cpu, + forward_batch.seq_lens_sum, + prefix_lens, + prefill_wrappers=self.prefill_wrappers_paged, + use_ragged=use_ragged, + encoder_lens=forward_batch.encoder_lens, + spec_info=None, + ) + self.forward_metadata = PrefillMetadata( + self.prefill_wrappers_paged, use_ragged, extend_no_prefix + ) + + def init_cuda_graph_state( + self, + max_bs: int, + max_num_tokens: int, + kv_indices_buf: Optional[torch.Tensor] = None, + ): + if kv_indices_buf is None: + cuda_graph_kv_indices = torch.zeros( + (max_num_tokens * self.max_context_len,), + dtype=torch.int32, + device="cuda", + ) + else: + cuda_graph_kv_indices = kv_indices_buf + + self.cuda_graph_kv_indices = [cuda_graph_kv_indices] + [ + cuda_graph_kv_indices.clone() for _ in range(self.num_wrappers - 1) + ] + + # Ensure tensors are properly allocated + for i in range(self.num_wrappers): + # Force allocation by performing a small operation + if len(self.cuda_graph_kv_indices[i]) > 0: + self.cuda_graph_kv_indices[i][0] = 0 + + if not self.skip_prefill: + self.cuda_graph_custom_mask = torch.zeros( + (max_num_tokens * self.max_context_len), + dtype=torch.uint8, + device="cuda", + ) + self.cuda_graph_qk_indptr = [x.clone() for x in self.kv_indptr] + self.cuda_graph_qo_indptr = [x.clone() for x in self.kv_indptr] + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + if forward_mode.is_decode_or_idle(): + decode_wrappers = [] + for i in range(self.num_wrappers): + decode_wrappers.append( + BatchDecodeWithPagedKVCacheWrapper( + self.workspace_buffer, + "NHD", + use_cuda_graph=True, + use_tensor_cores=self.decode_use_tensor_cores, + paged_kv_indptr_buffer=self.kv_indptr[i][: num_tokens + 1], + paged_kv_indices_buffer=self.cuda_graph_kv_indices[i], + paged_kv_last_page_len_buffer=self.kv_last_page_len[ + :num_tokens + ], + ) + ) + seq_lens_sum = seq_lens.sum().item() + self.indices_updater_decode.update( + req_pool_indices, + seq_lens, + seq_lens.cpu(), # may add a little overhead in capture stage + seq_lens_sum, + decode_wrappers=decode_wrappers, + encoder_lens=encoder_lens, + spec_info=spec_info, + ) + self.decode_cuda_graph_metadata[bs] = decode_wrappers + self.forward_metadata = DecodeMetadata(decode_wrappers) + for i in range(self.num_wrappers): + decode_wrappers[i].begin_forward = partial( + fast_decode_plan, decode_wrappers[i] + ) + elif forward_mode.is_target_verify(): + prefill_wrappers = [] + for i in range(self.num_wrappers): + prefill_wrappers.append( + BatchPrefillWithPagedKVCacheWrapper( + self.workspace_buffer, + "NHD", + use_cuda_graph=True, + qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1], + paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1], + paged_kv_indices_buf=self.cuda_graph_kv_indices[i], + paged_kv_last_page_len_buf=self.kv_last_page_len[:bs], + custom_mask_buf=self.cuda_graph_custom_mask, + mask_indptr_buf=self.cuda_graph_qk_indptr[i][: bs + 1], + ) + ) + seq_lens_sum = seq_lens.sum().item() + self.indices_updater_prefill.update( + req_pool_indices, + seq_lens, + seq_lens.cpu(), # may add a little overhead in capture stage + seq_lens_sum, + prefix_lens=None, + prefill_wrappers=prefill_wrappers, + use_ragged=False, + encoder_lens=encoder_lens, + spec_info=spec_info, + ) + self.prefill_cuda_graph_metadata[bs] = prefill_wrappers + self.forward_metadata = PrefillMetadata(prefill_wrappers, False, False) + elif forward_mode.is_draft_extend(): + prefill_wrappers = [] + for i in range(self.num_wrappers): + prefill_wrappers.append( + BatchPrefillWithPagedKVCacheWrapper( + self.workspace_buffer, + "NHD", + backend="fa2", + use_cuda_graph=True, + qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1], + paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1], + paged_kv_indices_buf=self.cuda_graph_kv_indices[i], + paged_kv_last_page_len_buf=self.kv_last_page_len[:bs], + ) + ) + + seq_lens_sum = seq_lens.sum().item() + self.indices_updater_prefill.update( + req_pool_indices, + seq_lens, + seq_lens.cpu(), # may add a little overhead in capture stage + seq_lens_sum, + prefix_lens=None, + prefill_wrappers=prefill_wrappers, + use_ragged=False, + encoder_lens=encoder_lens, + spec_info=spec_info, + ) + self.prefill_cuda_graph_metadata[bs] = prefill_wrappers + self.forward_metadata = PrefillMetadata(prefill_wrappers, False, False) + else: + raise ValueError(f"Invalid mode: {forward_mode=}") + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + ): + if forward_mode.is_decode_or_idle(): + self.indices_updater_decode.update( + req_pool_indices[:bs], + seq_lens[:bs], + seq_lens_cpu[:bs] if seq_lens_cpu is not None else None, + seq_lens_sum, + decode_wrappers=self.decode_cuda_graph_metadata[bs], + encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None, + spec_info=spec_info, + ) + elif forward_mode.is_target_verify(): + self.indices_updater_prefill.update( + req_pool_indices[:bs], + seq_lens[:bs], + seq_lens_cpu[:bs] if seq_lens_cpu is not None else None, + seq_lens_sum, + prefix_lens=None, + prefill_wrappers=self.prefill_cuda_graph_metadata[bs], + use_ragged=False, + encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None, + spec_info=spec_info, + ) + elif forward_mode.is_draft_extend(): + self.indices_updater_prefill.update( + req_pool_indices[:bs], + seq_lens[:bs], + seq_lens_cpu[:bs] if seq_lens_cpu is not None else None, + seq_lens_sum, + prefix_lens=None, + prefill_wrappers=self.prefill_cuda_graph_metadata[bs], + use_ragged=False, + encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None, + spec_info=spec_info, + ) + else: + raise ValueError("Invalid forward mode") + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + prefill_wrapper_paged = self.forward_metadata.prefill_wrappers[ + self._get_wrapper_idx(layer) + ] + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + + logits_soft_cap = layer.logit_cap + + q = q.contiguous() + if not self.forward_metadata.use_ragged: + if k is not None: + assert v is not None + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, layer.k_scale, layer.v_scale + ) + + o = prefill_wrapper_paged.forward( + q.view(-1, layer.tp_q_head_num, layer.head_dim), + forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id), + causal=not layer.is_cross_attention, + sm_scale=layer.scaling, + window_left=layer.sliding_window_size, + logits_soft_cap=logits_soft_cap, + # Must use _float to avoid device-to-host copy that breaks cuda graph capture. + k_scale=layer.k_scale_float, + v_scale=layer.v_scale_float, + ) + else: + causal = True + if layer.attn_type == AttentionType.ENCODER_ONLY: + save_kv_cache = False + causal = False + + if self.forward_metadata.extend_no_prefix: + # NOTE: FlashInfer currently has limitations with head_dim = 32 or other dimensions + # The FlashInfer head_dim limitation itself is tracked here: + # https://github.com/flashinfer-ai/flashinfer/issues/1048 + o = self.prefill_wrapper_ragged.forward( + q.view(-1, layer.tp_q_head_num, layer.head_dim), + k.view(-1, layer.tp_k_head_num, layer.head_dim), + v.view(-1, layer.tp_v_head_num, layer.head_dim), + causal=causal, + sm_scale=layer.scaling, + logits_soft_cap=logits_soft_cap, + ) + + else: + o1, s1 = self.prefill_wrapper_ragged.forward_return_lse( + q.view(-1, layer.tp_q_head_num, layer.head_dim), + k.view(-1, layer.tp_k_head_num, layer.head_dim), + v.view(-1, layer.tp_v_head_num, layer.head_dim), + causal=True, + sm_scale=layer.scaling, + logits_soft_cap=logits_soft_cap, + ) + o2, s2 = prefill_wrapper_paged.forward_return_lse( + q.view(-1, layer.tp_q_head_num, layer.head_dim), + forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id), + causal=False, + sm_scale=layer.scaling, + logits_soft_cap=logits_soft_cap, + ) + + o, _ = merge_state(o1, s1, o2, s2) + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, layer.k_scale, layer.v_scale + ) + + return o.view(-1, layer.tp_q_head_num * layer.head_dim) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + decode_wrapper = self.forward_metadata.decode_wrappers[ + self._get_wrapper_idx(layer) + ] + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + + if k is not None: + assert v is not None + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, layer.k_scale, layer.v_scale + ) + + # Call the wrapped function + o = decode_wrapper.forward( + q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id), + sm_scale=layer.scaling, + logits_soft_cap=layer.logit_cap, + # Must use _float to avoid device-to-host copy that breaks cuda graph capture. + k_scale=layer.k_scale_float, + v_scale=layer.v_scale_float, + ) + + return o.view(-1, layer.tp_q_head_num * layer.head_dim) + + def _get_wrapper_idx(self, layer: RadixAttention): + if self.num_wrappers == 1: + return 0 + + if self.dispatch_reason == WrapperDispatch.SLIDING_WINDOW: + return layer.sliding_window_size == -1 + if self.dispatch_reason == WrapperDispatch.CROSS_ATTENTION: + return layer.is_cross_attention + + raise ValueError(f"Unknown dispatch reason: {self.dispatch_reason}") + + +class FlashInferIndicesUpdaterDecode: + def __init__(self, model_runner: ModelRunner, attn_backend: FlashInferAttnBackend): + # Parse Constants + self.num_qo_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.num_kv_heads = model_runner.model_config.get_num_kv_heads( + get_attention_tp_size() + ) + self.head_dim = model_runner.model_config.head_dim + self.data_type = model_runner.kv_cache_dtype + self.q_data_type = model_runner.dtype + self.sliding_window_size = model_runner.sliding_window_size + self.attn_backend = attn_backend + + # Buffers and wrappers + self.kv_indptr = attn_backend.kv_indptr + self.kv_last_page_len = attn_backend.kv_last_page_len + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.token_to_kv_pool_allocator = model_runner.token_to_kv_pool_allocator + + # Dispatch the update function + if self.attn_backend.dispatch_reason == WrapperDispatch.SLIDING_WINDOW: + self.update = self.update_sliding_window + elif self.attn_backend.dispatch_reason == WrapperDispatch.CROSS_ATTENTION: + self.update = self.update_cross_attention + else: + assert self.attn_backend.num_wrappers == 1 + self.update = self.update_single_wrapper + + def update( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + seq_lens_sum: int, + decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper], + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + # Keep the signature for type checking. It will be assigned during runtime. + raise NotImplementedError() + + def update_single_wrapper( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + seq_lens_sum: int, + decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper], + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + decode_wrappers = decode_wrappers or self.decode_wrappers + self.call_begin_forward( + decode_wrappers[0], + req_pool_indices, + seq_lens, + seq_lens_sum, + self.kv_indptr[0], + None, + spec_info, + seq_lens_cpu, + ) + + def update_sliding_window( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + seq_lens_sum: int, + decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper], + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + assert self.sliding_window_size is not None + for wrapper_id in range(2): + if wrapper_id == 0: + # Sliding window attention + paged_kernel_lens_tmp = torch.clamp( + seq_lens, max=self.sliding_window_size + 1 + ) + if seq_lens_cpu is not None: + seq_lens_cpu_tmp = torch.clamp( + seq_lens_cpu, max=self.sliding_window_size + 1 + ) + paged_kernel_lens_sum_tmp = seq_lens_cpu_tmp.sum().item() + else: + paged_kernel_lens_sum_tmp = paged_kernel_lens_tmp.sum().item() + kv_start_idx_tmp = seq_lens - paged_kernel_lens_tmp + else: + # Full attention + paged_kernel_lens_tmp = seq_lens + paged_kernel_lens_sum_tmp = seq_lens_sum + seq_lens_cpu_tmp = seq_lens_cpu + kv_start_idx_tmp = None + + use_sliding_window_kv_pool = wrapper_id == 0 and isinstance( + self.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator + ) + + self.call_begin_forward( + decode_wrappers[wrapper_id], + req_pool_indices, + paged_kernel_lens_tmp, + paged_kernel_lens_sum_tmp, + self.kv_indptr[wrapper_id], + kv_start_idx_tmp, + spec_info, + seq_lens_cpu=seq_lens_cpu_tmp, + use_sliding_window_kv_pool=use_sliding_window_kv_pool, + ) + + def update_cross_attention( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + seq_lens_sum: int, + decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper], + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + for wrapper_id in range(2): + if wrapper_id == 0: + # Normal attention + paged_kernel_lens = seq_lens + kv_start_idx = encoder_lens + else: + # Cross attention + paged_kernel_lens = encoder_lens + kv_start_idx = torch.zeros_like(encoder_lens) + seq_lens_sum = encoder_lens.sum().item() + + self.call_begin_forward( + decode_wrappers[wrapper_id], + req_pool_indices, + paged_kernel_lens, + seq_lens_sum, + self.kv_indptr[wrapper_id], + kv_start_idx, + spec_info, + seq_lens_cpu=seq_lens_cpu, + ) + + def call_begin_forward( + self, + wrapper: BatchDecodeWithPagedKVCacheWrapper, + req_pool_indices: torch.Tensor, + paged_kernel_lens: torch.Tensor, + paged_kernel_lens_sum: int, + kv_indptr: torch.Tensor, + kv_start_idx: torch.Tensor, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + use_sliding_window_kv_pool: bool = False, + ): + if spec_info is None: + bs = len(req_pool_indices) + kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + + if wrapper.is_cuda_graph_enabled: + # Directly write to the cuda graph input buffer + kv_indices = wrapper._paged_kv_indices_buf + else: + kv_indices = torch.empty( + paged_kernel_lens_sum, dtype=torch.int32, device="cuda" + ) + + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + paged_kernel_lens, + kv_indptr, + kv_start_idx, + kv_indices, + self.req_to_token.shape[1], + ) + else: + kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices + bs = kv_indptr.shape[0] - 1 + + if use_sliding_window_kv_pool: + kv_last_index = kv_indptr[-1] + kv_indices[:kv_last_index] = ( + self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( + kv_indices[:kv_last_index] + ) + ) + + global global_override_indptr_cpu + locally_override = False + if seq_lens_cpu is not None and global_override_indptr_cpu is None: + locally_override = True + global_override_indptr_cpu = torch.empty_like(kv_indptr, device="cpu") + global_override_indptr_cpu[0] = 0 + global_override_indptr_cpu[1 : bs + 1] = torch.cumsum(seq_lens_cpu, dim=0) + + wrapper.begin_forward( + kv_indptr, + kv_indices, + self.kv_last_page_len[:bs], + self.num_qo_heads, + self.num_kv_heads, + self.head_dim, + 1, + data_type=self.data_type, + q_data_type=self.q_data_type, + non_blocking=True, + ) + + if locally_override: + global_override_indptr_cpu = None + + +class FlashInferIndicesUpdaterPrefill: + def __init__(self, model_runner: ModelRunner, attn_backend: FlashInferAttnBackend): + # Parse Constants + self.num_qo_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.num_kv_heads = model_runner.model_config.get_num_kv_heads( + get_attention_tp_size() + ) + self.head_dim = model_runner.model_config.head_dim + self.data_type = model_runner.kv_cache_dtype + self.q_data_type = model_runner.dtype + self.sliding_window_size = model_runner.sliding_window_size + self.attn_backend = attn_backend + + # Buffers and wrappers + self.kv_indptr = attn_backend.kv_indptr + self.kv_last_page_len = attn_backend.kv_last_page_len + self.qo_indptr = attn_backend.qo_indptr + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.token_to_kv_pool_allocator = model_runner.token_to_kv_pool_allocator + self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged + + # Dispatch the update function + if self.attn_backend.dispatch_reason == WrapperDispatch.SLIDING_WINDOW: + self.update = self.update_sliding_window + elif self.attn_backend.dispatch_reason == WrapperDispatch.CROSS_ATTENTION: + self.update = self.update_cross_attention + else: + assert self.attn_backend.num_wrappers == 1 + self.update = self.update_single_wrapper + + def update( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + seq_lens_sum: int, + prefix_lens: torch.Tensor, + prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper], + use_ragged: bool, + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + # Keep the signature for type checking. It will be assigned during runtime. + raise NotImplementedError() + + def update_single_wrapper( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + seq_lens_sum: int, + prefix_lens: torch.Tensor, + prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper], + use_ragged: bool, + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + if use_ragged: + # TODO: remove this device sync, we can use forward_batch.extend_prefix_lens_cpu + # and forward_batch.extend_seq_lens_cpu + paged_kernel_lens = prefix_lens + paged_kernel_lens_sum = paged_kernel_lens.sum().item() + else: + paged_kernel_lens = seq_lens + paged_kernel_lens_sum = seq_lens_sum + + self.call_begin_forward( + self.prefill_wrapper_ragged, + prefill_wrappers[0], + req_pool_indices, + paged_kernel_lens, + paged_kernel_lens_sum, + seq_lens, + prefix_lens, + None, + self.kv_indptr[0], + self.qo_indptr[0], + use_ragged, + spec_info, + ) + + def update_sliding_window( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + seq_lens_sum: int, + prefix_lens: torch.Tensor, + prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper], + use_ragged: bool, + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + for wrapper_id in range(2): + if wrapper_id == 0: + # window attention use paged only + paged_kernel_lens = torch.minimum( + seq_lens, + torch.tensor(self.sliding_window_size) + seq_lens - prefix_lens, + ) + paged_kernel_lens_sum = paged_kernel_lens.sum().item() + else: + # full attention + paged_kernel_lens = seq_lens + paged_kernel_lens_sum = seq_lens_sum + + kv_start_idx = seq_lens - paged_kernel_lens + use_sliding_window_kv_pool = wrapper_id == 0 and isinstance( + self.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator + ) + + self.call_begin_forward( + self.prefill_wrapper_ragged, + prefill_wrappers[wrapper_id], + req_pool_indices, + paged_kernel_lens, + paged_kernel_lens_sum, + seq_lens, + prefix_lens, + kv_start_idx, + self.kv_indptr[wrapper_id], + self.qo_indptr[wrapper_id], + use_ragged, + spec_info, + use_sliding_window_kv_pool=use_sliding_window_kv_pool, + ) + + def update_cross_attention( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + seq_lens_sum: int, + prefix_lens: torch.Tensor, + prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper], + use_ragged: bool, + encoder_lens: Optional[torch.Tensor], + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + for wrapper_id in range(2): + if wrapper_id == 0: + # normal attention + paged_kernel_lens = seq_lens + kv_start_idx = encoder_lens + paged_kernel_lens_sum = seq_lens_sum + else: + # cross attention + paged_kernel_lens = encoder_lens + kv_start_idx = torch.zeros_like(encoder_lens) + paged_kernel_lens_sum = paged_kernel_lens.sum().item() + + self.call_begin_forward( + self.prefill_wrapper_ragged, + prefill_wrappers[wrapper_id], + req_pool_indices, + paged_kernel_lens, + paged_kernel_lens_sum, + seq_lens, + prefix_lens, + kv_start_idx, + self.kv_indptr[wrapper_id], + self.qo_indptr[wrapper_id], + use_ragged, + spec_info, + ) + + def call_begin_forward( + self, + wrapper_ragged: BatchPrefillWithRaggedKVCacheWrapper, + wrapper_paged: BatchPrefillWithPagedKVCacheWrapper, + req_pool_indices: torch.Tensor, + paged_kernel_lens: torch.Tensor, + paged_kernel_lens_sum: int, + seq_lens: torch.Tensor, + prefix_lens: torch.Tensor, + kv_start_idx: torch.Tensor, + kv_indptr: torch.Tensor, + qo_indptr: torch.Tensor, + use_ragged: bool, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + use_sliding_window_kv_pool: bool = False, + ): + bs = len(seq_lens) + if spec_info is None: + assert len(seq_lens) == len(req_pool_indices) + # Normal extend + kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + paged_kernel_lens_sum + 256, + dtype=torch.int32, + device=req_pool_indices.device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + paged_kernel_lens, + kv_indptr, + kv_start_idx, + kv_indices, + self.req_to_token.shape[1], + ) + qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0) + qo_indptr = qo_indptr[: bs + 1] + custom_mask = None + else: + assert isinstance(spec_info, EagleDraftInput) or isinstance( + spec_info, EagleVerifyInput + ) + kv_indices, kv_indptr, qo_indptr, custom_mask = ( + spec_info.generate_attn_arg_prefill( + req_pool_indices, + paged_kernel_lens, + paged_kernel_lens_sum, + self.req_to_token, + ) + ) + + # extend part + if use_ragged: + wrapper_ragged.begin_forward( + qo_indptr, + qo_indptr, + self.num_qo_heads, + self.num_kv_heads, + self.head_dim, + q_data_type=self.q_data_type, + ) + + if use_sliding_window_kv_pool: + kv_last_index = kv_indptr[-1] + kv_indices[:kv_last_index] = ( + self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( + kv_indices[:kv_last_index] + ) + ) + + # cached part + wrapper_paged.begin_forward( + qo_indptr, + kv_indptr, + kv_indices, + self.kv_last_page_len[:bs], + self.num_qo_heads, + self.num_kv_heads, + self.head_dim, + 1, + q_data_type=self.q_data_type, + kv_data_type=self.data_type, + custom_mask=custom_mask, + non_blocking=True, + ) + + +class FlashInferMultiStepDraftBackend: + """ + Wrap multiple flashinfer attention backends as one for multiple consecutive + draft decoding steps. + """ + + def __init__( + self, + model_runner: ModelRunner, + topk: int, + speculative_num_steps: int, + ): + from sglang.srt.speculative.eagle_utils import generate_draft_decode_kv_indices + + self.topk = topk + self.speculative_num_steps = speculative_num_steps + self.generate_draft_decode_kv_indices = generate_draft_decode_kv_indices + self.page_size = model_runner.page_size + + max_bs = model_runner.req_to_token_pool.size * self.topk + self.kv_indptr = torch.zeros( + ( + self.speculative_num_steps, + max_bs + 1, + ), + dtype=torch.int32, + device=model_runner.device, + ) + self.kv_last_page_len = torch.ones( + (max_bs,), dtype=torch.int32, device=model_runner.device + ) + self.attn_backends: List[FlashInferAttnBackend] = [] + for i in range(self.speculative_num_steps): + self.attn_backends.append( + FlashInferAttnBackend( + model_runner, + skip_prefill=True, + kv_indptr_buf=self.kv_indptr[i], + kv_last_page_len_buf=self.kv_last_page_len, + ) + ) + + self.max_context_len = self.attn_backends[0].max_context_len + + # Cached variables for generate_draft_decode_kv_indices + self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1] + + def common_template( + self, + forward_batch: ForwardBatch, + kv_indices_buffer: torch.Tensor, + call_fn: Callable, + ): + num_seqs = forward_batch.batch_size + bs = self.topk * num_seqs + seq_lens_sum = forward_batch.seq_lens_sum + + self.generate_draft_decode_kv_indices[ + (self.speculative_num_steps, num_seqs, self.topk) + ]( + forward_batch.req_pool_indices, + forward_batch.req_to_token_pool.req_to_token, + forward_batch.seq_lens, + kv_indices_buffer, + self.kv_indptr, + forward_batch.positions, + self.pool_len, + kv_indices_buffer.shape[1], + self.kv_indptr.shape[1], + next_power_of_2(num_seqs), + next_power_of_2(self.speculative_num_steps), + next_power_of_2(bs), + self.page_size, + ) + + assert forward_batch.spec_info is not None + assert isinstance(forward_batch.spec_info, EagleDraftInput) + + # Copy the kv_indptr once to avoid multiple device-to-host copies in flashinfer's plan. + indptr_cpu_whole = self.kv_indptr[:, : bs + 1].cpu() + global global_override_indptr_cpu + + for i in range(self.speculative_num_steps - 1): + forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1] + forward_batch.spec_info.kv_indices = kv_indices_buffer[i][ + : seq_lens_sum * self.topk + bs * (i + 1) + ] + global_override_indptr_cpu = indptr_cpu_whole[i] + call_fn(i, forward_batch) + + global_override_indptr_cpu = None + + def init_forward_metadata(self, forward_batch: ForwardBatch): + kv_indices = torch.empty( + ( + self.speculative_num_steps, + forward_batch.batch_size * self.topk * self.max_context_len, + ), + dtype=torch.int32, + device="cuda", + ) + + def call_fn(i, forward_batch): + forward_batch.spec_info.kv_indptr = ( + forward_batch.spec_info.kv_indptr.clone() + ) + forward_batch.spec_info.kv_indices = ( + forward_batch.spec_info.kv_indices.clone() + ) + self.attn_backends[i].init_forward_metadata(forward_batch) + + self.common_template(forward_batch, kv_indices, call_fn) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.cuda_graph_kv_indices = torch.zeros( + (self.speculative_num_steps, max_bs * self.max_context_len), + dtype=torch.int32, + device="cuda", + ) + + for i in range(self.speculative_num_steps): + self.attn_backends[i].init_cuda_graph_state( + max_bs, max_num_tokens, kv_indices_buf=self.cuda_graph_kv_indices[i] + ) + + def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_capture_cuda_graph( + forward_batch.batch_size, + forward_batch.batch_size * self.topk, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + ) + + self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn) + + def init_forward_metadata_replay_cuda_graph( + self, forward_batch: ForwardBatch, bs: int + ): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_replay_cuda_graph( + bs, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + seq_lens_sum=-1, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + seq_lens_cpu=forward_batch.seq_lens_cpu, + ) + + self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn) + + +def should_use_tensor_core( + kv_cache_dtype: torch.dtype, + num_attention_heads: int, + num_kv_heads: int, +) -> bool: + """ + Determine whether to use tensor cores for attention computation. + + Args: + kv_cache_dtype: Data type of the KV cache + num_attention_heads: Number of attention heads + num_kv_heads: Number of key/value heads + + Returns: + bool: Whether to use tensor cores + """ + # Try to use environment variable first + env_override = os.environ.get("SGLANG_FLASHINFER_USE_TENSOR_CORE") + if env_override is not None: + return env_override.lower() == "true" + + # Try to use _grouped_size_compiled_for_decode_kernels if available + # This is for flashinfer <=0.1.6. Otherwise, there is an accuracy bug + try: + from flashinfer.decode import _grouped_size_compiled_for_decode_kernels + + if not _grouped_size_compiled_for_decode_kernels( + num_attention_heads, + num_kv_heads, + ): + return True + else: + return False + except (ImportError, AttributeError): + pass + + # Calculate GQA group size + gqa_group_size = num_attention_heads // num_kv_heads + + # For Flashinfer, a GQA group size of at least 4 is needed to efficiently + # use Tensor Cores, as it fuses the head group with the token dimension in MMA. + if kv_cache_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + return True + elif kv_cache_dtype in (torch.float16, torch.half, torch.bfloat16): + return gqa_group_size >= 4 + else: + return False + + +# Use as a fast path to override the indptr in flashinfer's plan function +# This is used to remove some host-to-device copy overhead. +global_override_indptr_cpu = None + + +def fast_decode_plan( + self, + indptr: torch.Tensor, + indices: torch.Tensor, + last_page_len: torch.Tensor, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + page_size: int, + pos_encoding_mode: str = "NONE", + window_left: int = -1, + logits_soft_cap: Optional[float] = None, + q_data_type: Optional[Union[str, torch.dtype]] = None, + kv_data_type: Optional[Union[str, torch.dtype]] = None, + data_type: Optional[Union[str, torch.dtype]] = None, + sm_scale: Optional[float] = None, + rope_scale: Optional[float] = None, + rope_theta: Optional[float] = None, + non_blocking: bool = True, +) -> None: + """ + A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for FlashInferMultiStepDraftBackend. + Modifications: + - Remove unnecessary device-to-device copy for the cuda graph buffers. + - Remove unnecessary host-to-device copy for the metadata buffers. + """ + batch_size = len(last_page_len) + if logits_soft_cap is None: + logits_soft_cap = 0.0 + + # Handle data types consistently + if data_type is not None: + if q_data_type is None: + q_data_type = data_type + if kv_data_type is None: + kv_data_type = data_type + elif q_data_type is None: + q_data_type = "float16" + + if kv_data_type is None: + kv_data_type = q_data_type + + if self.use_tensor_cores: + qo_indptr_host = _get_range_buf(batch_size + 1, "cpu") + + if self.is_cuda_graph_enabled: + if batch_size != self._fixed_batch_size: + raise ValueError( + "The batch size should be fixed in cudagraph mode, the runtime batch size {} " + " mismatches the batch size set during initialization {}".format( + batch_size, self._fixed_batch_size + ) + ) + if len(indices) > len(self._paged_kv_indices_buf): + raise ValueError( + "The size of indices should be less than or equal to the allocated buffer" + ) + else: + self._paged_kv_indptr_buf = indptr + self._paged_kv_indices_buf = indices + self._paged_kv_last_page_len_buf = last_page_len + if self.use_tensor_cores: + self._qo_indptr_buf = qo_indptr_host.to( + self.device, non_blocking=non_blocking + ) + + # Create empty tensors for dtype info if needed + empty_q_data = torch.empty( + 0, + dtype=( + getattr(torch, q_data_type) if isinstance(q_data_type, str) else q_data_type + ), + device=self.device, + ) + + empty_kv_cache = torch.empty( + 0, + dtype=( + getattr(torch, kv_data_type) + if isinstance(kv_data_type, str) + else kv_data_type + ), + device=self.device, + ) + + indptr_host = ( + global_override_indptr_cpu + if global_override_indptr_cpu is not None + else indptr.cpu() + ) + + with torch.cuda.device(self.device): + + if self.use_tensor_cores: + # ALSO convert last_page_len to CPU + if page_size == 1: + # When page size is 1, last_page_len is always 1. + # Directly construct the host tensor rather than executing a device-to-host copy. + last_page_len_host = torch.ones( + (batch_size,), dtype=torch.int32, device="cpu" + ) + else: + last_page_len_host = last_page_len.cpu() + + kv_lens_arr_host = get_seq_lens(indptr_host, last_page_len_host, page_size) + + try: + # Make sure we pass exactly 15 arguments for tensor core version + self._plan_info = self._cached_module.plan( + self._float_workspace_buffer, + self._int_workspace_buffer, + self._pin_memory_int_workspace_buffer, + qo_indptr_host, + indptr_host, + kv_lens_arr_host, + batch_size, # total_num_rows + batch_size, + num_qo_heads, + num_kv_heads, + page_size, + self.is_cuda_graph_enabled, + head_dim, + head_dim, + False, # causal + ) + except Exception as e: + raise RuntimeError(f"Error in standard plan: {e}") + else: + try: + # Make sure we pass exactly 15 arguments for standard version + self._plan_info = self._cached_module.plan( + self._float_workspace_buffer, + self._int_workspace_buffer, + self._pin_memory_int_workspace_buffer, + indptr_host, + batch_size, + num_qo_heads, + num_kv_heads, + page_size, + self.is_cuda_graph_enabled, + window_left, + logits_soft_cap, + head_dim, + head_dim, + empty_q_data, + empty_kv_cache, + ) + except Exception as e: + raise RuntimeError(f"Error in standard plan: {e}") + + self._pos_encoding_mode = pos_encoding_mode + self._window_left = window_left + self._logits_soft_cap = logits_soft_cap + self._sm_scale = sm_scale + self._rope_scale = rope_scale + self._rope_theta = rope_theta diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py new file mode 100644 index 000000000..05e9bef80 --- /dev/null +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -0,0 +1,1079 @@ +from __future__ import annotations + +""" +Support attention backend for flashinfer MLA. +The flashinfer_mla_disable_ragged flag controls whether to use ragged prefill wrapper and defaults to be false. +When it's set to false, all wrappers are BatchMLAPaged wrapper. +When it's set to true, the backend uses BatchRagged and BatchMLAPaged wrapper for prefilling, +and uses BatchMLAPaged wrapper for decoding. +More details can be found in https://docs.flashinfer.ai/api/mla.html +""" + +import os +from dataclasses import dataclass +from functools import partial +from typing import TYPE_CHECKING, Callable, Optional, Union + +import torch + +if os.environ["SGLANG_ENABLE_TORCH_COMPILE"] == "1": + import logging + + torch._logging.set_logs(dynamo=logging.ERROR) + torch._dynamo.config.suppress_errors = True + +from sglang.global_config import global_config +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.flashinfer_backend import ( + create_flashinfer_kv_indices_triton, +) +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput +from sglang.srt.utils import ( + is_flashinfer_available, + is_sm100_supported, + next_power_of_2, +) + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpecInfo + +if is_flashinfer_available(): + from flashinfer import ( + BatchMLAPagedAttentionWrapper, + BatchPrefillWithRaggedKVCacheWrapper, + ) + + +@dataclass +class DecodeMetadata: + decode_wrapper: BatchMLAPagedAttentionWrapper + + +@dataclass +class PrefillMetadata: + prefill_wrapper: BatchMLAPagedAttentionWrapper + use_ragged: bool + + +# Reuse this workspace buffer across all flashinfer wrappers +global_workspace_buffer = None + + +class FlashInferMhaChunkKVRunner: + def __init__( + self, model_runner: ModelRunner, attn_backend: "FlashInferMlaAttnBackend" + ): + # Parse Constants + self.num_local_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim + self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim + self.v_head_dim = model_runner.model_config.v_head_dim + self.data_type = model_runner.dtype + self.q_data_type = model_runner.dtype + + # Buffers and wrappers + self.qo_indptr = attn_backend.qo_indptr + self.workspace_buffer = attn_backend.workspace_buffer + self.fmha_backend = attn_backend.fmha_backend + + self.chunk_ragged_wrappers = [] + self.ragged_wrapper = attn_backend.prefill_wrapper_ragged + + def update_prefix_chunks(self, num_prefix_chunks: int): + while num_prefix_chunks > len(self.chunk_ragged_wrappers): + ragged_wrapper = BatchPrefillWithRaggedKVCacheWrapper( + self.workspace_buffer, "NHD", backend=self.fmha_backend + ) + self.chunk_ragged_wrappers.append(ragged_wrapper) + + def update_wrapper( + self, + forward_batch: ForwardBatch, + disable_flashinfer_ragged: bool = False, + ): + assert forward_batch.num_prefix_chunks is not None + num_prefix_chunks = forward_batch.num_prefix_chunks + self.update_prefix_chunks(num_prefix_chunks) + + prefix_lens = forward_batch.extend_prefix_lens + seq_lens = forward_batch.seq_lens + + bs = len(seq_lens) + qo_indptr = self.qo_indptr + qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0) + qo_indptr = qo_indptr[: bs + 1] + + for chunk_idx in range(forward_batch.num_prefix_chunks): + # MHA for chunked prefix kv cache when running model with MLA + assert forward_batch.prefix_chunk_idx is not None + assert forward_batch.prefix_chunk_cu_seq_lens is not None + assert forward_batch.prefix_chunk_max_seq_lens is not None + + kv_indptr = forward_batch.prefix_chunk_cu_seq_lens[chunk_idx] + wrapper = self.chunk_ragged_wrappers[chunk_idx] + wrapper.begin_forward( + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + num_qo_heads=self.num_local_heads, + num_kv_heads=self.num_local_heads, + head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim, + head_dim_vo=self.v_head_dim, + q_data_type=self.q_data_type, + causal=False, + ) + # ragged prefill + if not disable_flashinfer_ragged: + self.ragged_wrapper.begin_forward( + qo_indptr=qo_indptr, + kv_indptr=qo_indptr, + num_qo_heads=self.num_local_heads, + num_kv_heads=self.num_local_heads, + head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim, + head_dim_vo=self.v_head_dim, + q_data_type=self.q_data_type, + causal=True, + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + ): + logits_soft_cap = layer.logit_cap + if forward_batch.attn_attend_prefix_cache: + chunk_idx = forward_batch.prefix_chunk_idx + assert chunk_idx >= 0 + wrapper = self.chunk_ragged_wrappers[chunk_idx] + o1, s1 = wrapper.forward_return_lse( + q.view(-1, layer.tp_q_head_num, layer.head_dim), + k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype), + v.view(-1, layer.tp_v_head_num, layer.v_head_dim).to(q.dtype), + causal=False, + sm_scale=layer.scaling, + logits_soft_cap=logits_soft_cap, + ) + else: + o1, s1 = self.ragged_wrapper.forward_return_lse( + q.view(-1, layer.tp_q_head_num, layer.head_dim), + k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype), + v.view(-1, layer.tp_v_head_num, layer.v_head_dim).to(q.dtype), + causal=True, + sm_scale=layer.scaling, + logits_soft_cap=logits_soft_cap, + ) + + return o1, s1 + + +class FlashInferMLAAttnBackend(AttentionBackend): + """Flashinfer attention kernels.""" + + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + kv_indptr_buf: Optional[torch.Tensor] = None, + q_indptr_decode_buf: Optional[torch.Tensor] = None, + ): + super().__init__() + + # Parse constants + self.max_context_len = model_runner.model_config.context_len + self.device = model_runner.device + self.skip_prefill = skip_prefill + self.enable_chunk_kv = ( + not skip_prefill + and global_server_args_dict["disaggregation_mode"] != "decode" + and not global_server_args_dict["disable_chunked_prefix_cache"] + and not global_server_args_dict["flashinfer_mla_disable_ragged"] + ) + self.page_size = model_runner.page_size + + # Allocate buffers + global global_workspace_buffer + if global_workspace_buffer is None: + # different from flashinfer zero_init_global_workspace_buffer + global_workspace_buffer = torch.empty( + global_config.flashinfer_workspace_size, + dtype=torch.uint8, + device=model_runner.device, + ) + self.workspace_buffer = global_workspace_buffer + + max_bs = model_runner.req_to_token_pool.size + if kv_indptr_buf is None: + self.kv_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + else: + self.kv_indptr = kv_indptr_buf + + if not self.skip_prefill: + self.qo_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + + if q_indptr_decode_buf is None: + self.q_indptr_decode = torch.arange( + 0, max_bs + 1, dtype=torch.int32, device=model_runner.device + ) + else: + self.q_indptr_decode = q_indptr_decode_buf + + self.fmha_backend = "auto" + if is_sm100_supported(): + self.fmha_backend = "cutlass" + self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper( + self.workspace_buffer, "NHD", backend=self.fmha_backend + ) + + if not self.skip_prefill: + self.prefill_wrapper_paged = BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + backend="auto", + ) + + # FlashinferMLA backend uses mla wrapper for target verify + self.prefill_wrapper_verify = BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + backend="auto", + ) + + self.decode_wrapper = BatchMLAPagedAttentionWrapper( + self.workspace_buffer, backend="auto" + ) + + # Create indices updater + if not skip_prefill: + self.indices_updater_prefill = FlashInferMLAIndicesUpdaterPrefill( + model_runner, self + ) + if self.enable_chunk_kv: + self.mha_chunk_kv_cache = FlashInferMhaChunkKVRunner(model_runner, self) + + self.indices_updater_decode = FlashInferMLAIndicesUpdaterDecode( + model_runner, self + ) + + # Other metadata + self.forward_metadata: Union[PrefillMetadata, DecodeMetadata] = None + self.decode_cuda_graph_metadata = {} + self.prefill_cuda_graph_metadata = {} # For verify + + def init_forward_metadata(self, forward_batch: ForwardBatch): + if forward_batch.forward_mode.is_decode_or_idle(): + self.indices_updater_decode.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + decode_wrapper=self.decode_wrapper, + init_metadata_replay=False, + ) + self.forward_metadata = DecodeMetadata(self.decode_wrapper) + elif forward_batch.forward_mode.is_draft_extend(): + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + prefix_lens=None, + prefill_wrapper_paged=self.prefill_wrapper_paged, + use_ragged=False, + spec_info=forward_batch.spec_info, + ) + self.forward_metadata = PrefillMetadata(self.prefill_wrapper_paged, False) + elif forward_batch.forward_mode.is_target_verify(): + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + prefix_lens=None, + prefill_wrapper_paged=self.prefill_wrapper_verify, + use_ragged=False, + spec_info=forward_batch.spec_info, + ) + self.forward_metadata = PrefillMetadata(self.prefill_wrapper_verify, False) + else: + prefix_lens = forward_batch.extend_prefix_lens + extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu) + use_ragged = ( + not global_server_args_dict["flashinfer_mla_disable_ragged"] + and extend_no_prefix + ) + + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + prefix_lens, + prefill_wrapper_paged=self.prefill_wrapper_paged, + use_ragged=use_ragged, + ) + self.forward_metadata = PrefillMetadata( + self.prefill_wrapper_paged, use_ragged + ) + + def init_cuda_graph_state( + self, + max_bs: int, + max_num_tokens: int, + kv_indices_buf: Optional[torch.Tensor] = None, + ): + if kv_indices_buf is None: + cuda_graph_kv_indices = torch.zeros( + (max_bs * self.max_context_len,), + dtype=torch.int32, + device="cuda", + ) + else: + cuda_graph_kv_indices = kv_indices_buf + + self.cuda_graph_kv_indices = cuda_graph_kv_indices + self.cuda_graph_qo_indptr = self.q_indptr_decode.clone() + self.cuda_graph_kv_indptr = self.kv_indptr.clone() + self.cuda_graph_kv_lens = torch.ones( + (max_bs,), dtype=torch.int32, device=self.device + ) + + # For fast decode plan in graph replaying + self.cuda_graph_qo_indptr_cpu = self.cuda_graph_qo_indptr.to("cpu") + self.cuda_graph_kv_indptr_cpu = self.cuda_graph_kv_indptr.to("cpu") + self.fast_decode_kwargs = { + "qo_indptr_cpu": self.cuda_graph_qo_indptr_cpu, + "kv_indptr_cpu": self.cuda_graph_kv_indptr_cpu, + "kv_indices": self.cuda_graph_kv_indices, + } + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + ): + if forward_mode.is_decode_or_idle(): + decode_wrapper = BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + use_cuda_graph=True, + qo_indptr=self.cuda_graph_qo_indptr[: num_tokens + 1], + kv_indptr=self.cuda_graph_kv_indptr[: num_tokens + 1], + kv_indices=self.cuda_graph_kv_indices, + kv_len_arr=self.cuda_graph_kv_lens[:num_tokens], + backend="auto", + ) + + seq_lens_sum = seq_lens.sum().item() + self.indices_updater_decode.update( + req_pool_indices, + seq_lens, + seq_lens_sum, + decode_wrapper=decode_wrapper, + init_metadata_replay=False, + spec_info=spec_info, + ) + self.decode_cuda_graph_metadata[bs] = decode_wrapper + self.forward_metadata = DecodeMetadata(decode_wrapper) + decode_wrapper.plan = partial(fast_mla_decode_plan, decode_wrapper) + elif forward_mode.is_target_verify(): + verify_wrapper = BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + use_cuda_graph=True, + qo_indptr=self.cuda_graph_qo_indptr[: bs + 1], + kv_indptr=self.cuda_graph_kv_indptr[: bs + 1], + kv_indices=self.cuda_graph_kv_indices, + kv_len_arr=self.cuda_graph_kv_lens[:bs], + backend="auto", + ) + seq_lens_sum = seq_lens.sum().item() + self.indices_updater_prefill.update( + req_pool_indices, + seq_lens, + seq_lens_sum, + prefix_lens=None, + prefill_wrapper_paged=verify_wrapper, + use_ragged=False, + spec_info=spec_info, + ) + self.prefill_cuda_graph_metadata[bs] = verify_wrapper + self.forward_metadata = PrefillMetadata(verify_wrapper, False) + elif forward_mode.is_draft_extend(): + draft_extend_wrapper = BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + use_cuda_graph=True, + qo_indptr=self.cuda_graph_qo_indptr[: bs + 1], + kv_indptr=self.cuda_graph_kv_indptr[: bs + 1], + kv_indices=self.cuda_graph_kv_indices, + kv_len_arr=self.cuda_graph_kv_lens[:bs], + backend="auto", + ) + seq_lens_sum = seq_lens.sum().item() + self.indices_updater_prefill.update( + req_pool_indices, + seq_lens, + seq_lens_sum, + prefix_lens=None, + prefill_wrapper_paged=draft_extend_wrapper, + use_ragged=False, + spec_info=spec_info, + ) + self.prefill_cuda_graph_metadata[bs] = draft_extend_wrapper + self.forward_metadata = PrefillMetadata(draft_extend_wrapper, False) + else: + raise ValueError(f"Invalid mode: {forward_mode=}") + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + seq_lens_cpu: Optional[torch.Tensor], + ): + if forward_mode.is_decode_or_idle(): + assert seq_lens_cpu is not None + kv_len_arr_cpu = seq_lens_cpu[:bs] + self.cuda_graph_kv_indptr_cpu[1 : bs + 1] = torch.cumsum( + kv_len_arr_cpu, dim=0 + ) + self.fast_decode_kwargs.update( + { + "qo_indptr_cpu": self.cuda_graph_qo_indptr_cpu[: bs + 1], + "kv_indptr_cpu": self.cuda_graph_kv_indptr_cpu[: bs + 1], + "kv_len_arr_cpu": kv_len_arr_cpu, + } + ) + + self.indices_updater_decode.update( + req_pool_indices[:bs], + seq_lens[:bs], + seq_lens_sum, + decode_wrapper=self.decode_cuda_graph_metadata[bs], + init_metadata_replay=True, + spec_info=spec_info, + **self.fast_decode_kwargs, + ) + elif forward_mode.is_target_verify(): + self.indices_updater_prefill.update( + req_pool_indices[:bs], + seq_lens[:bs], + seq_lens_sum, + prefix_lens=None, + prefill_wrapper_paged=self.prefill_cuda_graph_metadata[bs], + use_ragged=False, + spec_info=spec_info, + ) + elif forward_mode.is_draft_extend(): + self.indices_updater_prefill.update( + req_pool_indices[:bs], + seq_lens[:bs], + seq_lens_sum, + prefix_lens=None, + prefill_wrapper_paged=self.prefill_cuda_graph_metadata[bs], + use_ragged=False, + spec_info=spec_info, + ) + else: + raise ValueError(f"Invalid forward mode: {forward_mode=}") + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + def init_mha_chunk_metadata( + self, forward_batch: ForwardBatch, disable_flashinfer_ragged: bool = False + ): + """Init the metadata for a forward pass.""" + self.mha_chunk_kv_cache.update_wrapper(forward_batch, disable_flashinfer_ragged) + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + ): + if ( + forward_batch.attn_attend_prefix_cache is not None + and forward_batch.mha_return_lse + ): # MHA Chunk + assert self.enable_chunk_kv + assert q_rope is None + assert k_rope is None + o1, s1 = self.mha_chunk_kv_cache.forward(q, k, v, layer, forward_batch) + return o1, s1 + + cache_loc = forward_batch.out_cache_loc + logits_soft_cap = layer.logit_cap + prefill_wrapper_paged = self.forward_metadata.prefill_wrapper + + # Save kv cache + if save_kv_cache and k is not None: + assert v is not None + if save_kv_cache: + if k_rope is not None: + forward_batch.token_to_kv_pool.set_mla_kv_buffer( + layer, cache_loc, k, k_rope + ) + else: + forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) + if q_rope is not None: + q = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + + if self.forward_metadata.use_ragged: + # ragged prefill + if q_rope is not None: + q = torch.cat([q, q_rope], dim=-1) + qall = q.view(-1, layer.tp_q_head_num, layer.head_dim) + if k_rope is not None: + k = torch.cat([k, k_rope], dim=-1) + o = self.prefill_wrapper_ragged.forward( + qall, + k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype), + v.view(-1, layer.tp_k_head_num, layer.v_head_dim).to(q.dtype), + causal=True, + sm_scale=layer.scaling, + logits_soft_cap=logits_soft_cap, + ) + else: + # mla paged prefill + k_buf = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id).to( + q.dtype + ) + if q_rope is None: + qall = q.view(-1, layer.tp_q_head_num, layer.head_dim) + q, q_rope = ( + qall[:, :, : layer.v_head_dim], + qall[:, :, layer.v_head_dim :], + ) + o = q.new_empty(q.shape) + o = prefill_wrapper_paged.run( + q, + q_rope, + k_buf[:, :, : layer.v_head_dim], + k_buf[:, :, layer.v_head_dim :], + out=o, + ) + + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + ): + decode_wrapper = self.forward_metadata.decode_wrapper + cache_loc = forward_batch.out_cache_loc + + if k is not None: + assert v is not None + if save_kv_cache: + if k_rope is not None: + forward_batch.token_to_kv_pool.set_mla_kv_buffer( + layer, + cache_loc, + k, + k_rope, + ) + else: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, + cache_loc, + k, + v, + ) + + # Reshape inputs + if q_rope is not None: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + else: + reshaped_q = q.view(-1, layer.tp_q_head_num, layer.head_dim) + q_nope = reshaped_q[:, :, : layer.v_head_dim] + q_rope = reshaped_q[:, :, layer.v_head_dim :] + + k_buffer = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id).to( + q.dtype + ) + + o = q_nope.new_empty(q_nope.shape) + # Direct call to run without the wrapper + o = decode_wrapper.run( + q_nope, + q_rope, + k_buffer[:, :, : layer.v_head_dim], + k_buffer[:, :, layer.v_head_dim :], + out=o, + ) + + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + +class FlashInferMLAIndicesUpdaterDecode: + def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend): + # Parse Constants + self.num_local_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.kv_lora_rank = model_runner.model_config.kv_lora_rank + self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim + self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim + self.scaling = model_runner.model_config.scaling + self.data_type = model_runner.dtype + self.attn_backend = attn_backend + + # Buffers and wrappers + self.kv_indptr = attn_backend.kv_indptr + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.q_indptr = attn_backend.q_indptr_decode + + def update( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + decode_wrapper: BatchMLAPagedAttentionWrapper, + init_metadata_replay: bool = False, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]] = None, + **fast_decode_kwargs, + ): + decode_wrapper = decode_wrapper or self.decode_wrapper + self.call_begin_forward( + decode_wrapper, + req_pool_indices, + seq_lens, + seq_lens_sum, + self.q_indptr, + self.kv_indptr, + init_metadata_replay, + spec_info, + **fast_decode_kwargs, + ) + + def call_begin_forward( + self, + wrapper: BatchMLAPagedAttentionWrapper, + req_pool_indices: torch.Tensor, + paged_kernel_lens: torch.Tensor, + paged_kernel_lens_sum: int, + q_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + init_metadata_replay: bool = False, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]] = None, + **fast_decode_kwargs, + ): + bs = len(req_pool_indices) + q_indptr = q_indptr[: bs + 1] + kv_lens = paged_kernel_lens.to(torch.int32) + sm_scale = self.scaling + if spec_info is None: + kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = ( + torch.empty(paged_kernel_lens_sum, dtype=torch.int32, device="cuda") + if not init_metadata_replay + else fast_decode_kwargs["kv_indices"] + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + paged_kernel_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.shape[1], + ) + else: + kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices + + if not init_metadata_replay: + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + self.num_local_heads, + self.kv_lora_rank, + self.qk_rope_head_dim, + 1, + False, + sm_scale, + self.data_type, + self.data_type, + ) + else: + wrapper.plan( + fast_decode_kwargs["qo_indptr_cpu"], + fast_decode_kwargs["kv_indptr_cpu"], + kv_indices, + fast_decode_kwargs["kv_len_arr_cpu"], + self.num_local_heads, + self.kv_lora_rank, + self.qk_rope_head_dim, + 1, + False, + sm_scale, + self.data_type, + self.data_type, + ) + + +class FlashInferMLAIndicesUpdaterPrefill: + def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend): + # Parse Constants + self.num_local_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.kv_lora_rank = model_runner.model_config.kv_lora_rank + self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim + self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim + self.v_head_dim = model_runner.model_config.v_head_dim + self.scaling = model_runner.model_config.scaling + self.data_type = model_runner.dtype + self.q_data_type = model_runner.dtype + self.attn_backend = attn_backend + + # Buffers and wrappers + self.kv_indptr = attn_backend.kv_indptr + self.qo_indptr = attn_backend.qo_indptr + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged + + def update( + self, + req_pool_indices: torch.Tnesor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + prefix_lens: torch.Tensor, + prefill_wrapper_paged: BatchMLAPagedAttentionWrapper, + use_ragged: bool, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]] = None, + ): + if use_ragged: + paged_kernel_lens = prefix_lens + paged_kernel_lens_sum = paged_kernel_lens.sum().item() + else: + paged_kernel_lens = seq_lens + paged_kernel_lens_sum = seq_lens_sum + + self.call_begin_forward( + self.prefill_wrapper_ragged, + prefill_wrapper_paged, + req_pool_indices, + paged_kernel_lens, + paged_kernel_lens_sum, + seq_lens, + prefix_lens, + self.kv_indptr, + self.qo_indptr, + use_ragged, + spec_info, + ) + + def call_begin_forward( + self, + wrapper_ragged: BatchPrefillWithRaggedKVCacheWrapper, + wrapper_paged: BatchMLAPagedAttentionWrapper, + req_pool_indices: torch.Tensor, + paged_kernel_lens: torch.Tensor, + paged_kernel_lens_sum: int, + seq_lens: torch.Tensor, + prefix_lens: torch.Tensor, + kv_indptr: torch.Tensor, + qo_indptr: torch.Tensor, + use_ragged: bool, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]] = None, + ): + bs = len(seq_lens) + sm_scale = self.scaling + + if spec_info is None: + assert len(seq_lens) == len(req_pool_indices) + kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + paged_kernel_lens_sum, + dtype=torch.int32, + device=req_pool_indices.device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + paged_kernel_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.shape[1], + ) + qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0) + qo_indptr = qo_indptr[: bs + 1] + custom_mask = None + else: + assert isinstance(spec_info, EagleDraftInput) or isinstance( + spec_info, EagleVerifyInput + ) + # TODO: Support topk > 1 with custom mask + kv_indices, kv_indptr, qo_indptr, custom_mask = ( + spec_info.generate_attn_arg_prefill( + req_pool_indices, + paged_kernel_lens, + paged_kernel_lens_sum, + self.req_to_token, + ) + ) + + if use_ragged: + # ragged prefill + wrapper_ragged.begin_forward( + qo_indptr=qo_indptr, + kv_indptr=qo_indptr, + num_qo_heads=self.num_local_heads, + num_kv_heads=self.num_local_heads, + head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim, + head_dim_vo=self.v_head_dim, + q_data_type=self.q_data_type, + causal=True, + ) + else: + # mla paged prefill + kv_len_arr = kv_indptr[1:] - kv_indptr[:-1] + wrapper_paged.plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + self.num_local_heads, + self.kv_lora_rank, + self.qk_rope_head_dim, + 1, + True, + sm_scale, + self.q_data_type, + self.data_type, + ) + + +class FlashInferMLAMultiStepDraftBackend: + """ + Wrap multiple flashinfer mla attention backends as one for multiple consecutive + draft decoding steps. + """ + + def __init__( + self, + model_runner: ModelRunner, + topk: int, + speculative_num_steps: int, + ): + from sglang.srt.speculative.eagle_utils import generate_draft_decode_kv_indices + + if topk > 1: + raise ValueError( + "Currently Flashinfer MLA only supports topk=1 for speculative decoding" + ) + self.topk = topk + self.speculative_num_steps = speculative_num_steps + self.generate_draft_decode_kv_indices = generate_draft_decode_kv_indices + + max_bs = model_runner.req_to_token_pool.size * self.topk + self.kv_indptr = torch.zeros( + ( + self.speculative_num_steps, + max_bs + 1, + ), + dtype=torch.int32, + device=model_runner.device, + ) + self.q_indptr_decode = torch.arange( + 0, max_bs + 1, dtype=torch.int32, device=model_runner.device + ) + + self.attn_backends = [] + for i in range(self.speculative_num_steps): + self.attn_backends.append( + FlashInferMLAAttnBackend( + model_runner, + skip_prefill=True, + kv_indptr_buf=self.kv_indptr[i], + q_indptr_decode_buf=self.q_indptr_decode, + ) + ) + + self.max_context_len = self.attn_backends[0].max_context_len + + # Cached variables for generate_draft_decode_kv_indices + self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1] + self.page_size = model_runner.server_args.page_size + + def common_template( + self, + forward_batch: ForwardBatch, + kv_indices_buffer: torch.Tensor, + call_fn: Callable, + ): + num_seqs = forward_batch.batch_size + bs = self.topk * num_seqs + seq_lens_sum = forward_batch.seq_lens_sum + + self.generate_draft_decode_kv_indices[ + (self.speculative_num_steps, num_seqs, self.topk) + ]( + forward_batch.req_pool_indices, + forward_batch.req_to_token_pool.req_to_token, + forward_batch.seq_lens, + kv_indices_buffer, + self.kv_indptr, + forward_batch.positions, + self.pool_len, + kv_indices_buffer.shape[1], + self.kv_indptr.shape[1], + next_power_of_2(num_seqs), + next_power_of_2(self.speculative_num_steps), + next_power_of_2(bs), + self.page_size, + ) + + assert forward_batch.spec_info is not None + assert isinstance(forward_batch.spec_info, EagleDraftInput) + + for i in range(self.speculative_num_steps - 1): + forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1] + forward_batch.spec_info.kv_indices = kv_indices_buffer[i][ + : seq_lens_sum * self.topk + bs * (i + 1) + ] + call_fn(i, forward_batch) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + kv_indices = torch.zeros( + ( + self.speculative_num_steps, + forward_batch.batch_size * self.topk * self.max_context_len, + ), + dtype=torch.int32, + device="cuda", + ) + + def call_fn(i, forward_batch): + assert forward_batch.spec_info is not None + assert isinstance(forward_batch.spec_info, EagleDraftInput) + forward_batch.spec_info.kv_indptr = ( + forward_batch.spec_info.kv_indptr.clone() + ) + forward_batch.spec_info.kv_indices = ( + forward_batch.spec_info.kv_indices.clone() + ) + self.attn_backends[i].init_forward_metadata(forward_batch) + + self.common_template(forward_batch, kv_indices, call_fn) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.cuda_graph_kv_indices = torch.zeros( + (self.speculative_num_steps, max_bs * self.max_context_len), + dtype=torch.int32, + device="cuda", + ) + + for i in range(self.speculative_num_steps): + self.attn_backends[i].init_cuda_graph_state( + max_bs, max_num_tokens, kv_indices_buf=self.cuda_graph_kv_indices[i] + ) + + def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_capture_cuda_graph( + forward_batch.batch_size, + forward_batch.batch_size * self.topk, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + ) + + self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn) + + def init_forward_metadata_replay_cuda_graph( + self, forward_batch: ForwardBatch, bs: int + ): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_replay_cuda_graph( + bs, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + seq_lens_sum=-1, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + seq_lens_cpu=forward_batch.seq_lens_cpu, + ) + + self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn) + + +def fast_mla_decode_plan( + self, + qo_indptr_cpu: torch.Tensor, + kv_indptr_cpu: torch.Tensor, + kv_indices: torch.Tensor, + kv_len_arr_cpu: torch.Tensor, + num_heads: int, + head_dim_ckv: int, + head_dim_kpe: int, + page_size: int, + causal: bool, + sm_scale: float, + q_data_type: torch.dtype, + kv_data_type: torch.dtype, +) -> None: + """A faster version of BatchMLAPagedAttentionWrapper::plan, + for skipping the stream synchronization in original plan function during + cuda graph replaying. + """ + self._causal = causal + self._page_size = page_size + self._sm_scale = sm_scale + + try: + # Standard version with just the required arguments (no use_profiler) + self._cached_module.plan.default( + self._float_workspace_buffer, + self._int_workspace_buffer, + self._pin_memory_int_workspace_buffer, + qo_indptr_cpu, + kv_indptr_cpu, + kv_len_arr_cpu, + num_heads, + head_dim_ckv, + causal, + ) + except Exception as e: + raise RuntimeError(f"Error in alternate MLA plan: {e}") diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py new file mode 100644 index 000000000..d1acb1a58 --- /dev/null +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -0,0 +1,541 @@ +from __future__ import annotations + +""" +Support attention backend for FlashMLA. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Optional, Tuple, Union + +import torch +import triton +from flash_mla import flash_mla_with_kvcache, get_mla_metadata + +from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend +from sglang.srt.layers.attention.utils import create_flashmla_kv_indices_triton +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpecInfo + + +# FlashMLA only supports pagesize=64 +PAGE_SIZE = 64 + +# FlashMLA FP8 issue: https://github.com/deepseek-ai/FlashMLA/issues/56 + + +@dataclass +class FlashMLADecodeMetadata: + flashmla_metadata: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + num_splits: Optional[torch.Tensor] = None + block_kv_indices: Optional[torch.Tensor] = None + + def __init__( + self, + flashmla_metadata: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + num_splits: Optional[torch.Tensor] = None, + block_kv_indices: Optional[torch.Tensor] = None, + ): + self.flashmla_metadata = flashmla_metadata + self.num_splits = num_splits + self.block_kv_indices = block_kv_indices + + +class FlashMLABackend(FlashInferMLAAttnBackend): + """Flashmla attention kernels.""" + + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + kv_indptr_buf: Optional[torch.Tensor] = None, + kv_last_page_len_buf: Optional[torch.Tensor] = None, + ): + super().__init__( + model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf + ) + + self.num_q_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.num_local_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.forward_metadata: Union[FlashMLADecodeMetadata] = None + self.kv_lora_rank = model_runner.model_config.kv_lora_rank + self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim + self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim + self.v_head_dim = model_runner.model_config.v_head_dim + self.scaling = model_runner.model_config.scaling + self.data_type = model_runner.kv_cache_dtype + self.q_data_type = model_runner.dtype + self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim + + self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens + + def init_forward_metadata(self, forward_batch: ForwardBatch): + + bs = forward_batch.batch_size + if forward_batch.forward_mode.is_decode_or_idle(): + max_seqlen_pad = triton.cdiv( + forward_batch.seq_lens_cpu.max().item(), PAGE_SIZE + ) + block_kv_indices = torch.full( + (bs, max_seqlen_pad), + -1, + dtype=torch.int32, + device=forward_batch.seq_lens.device, + ) + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + None, + block_kv_indices, + self.req_to_token.stride(0), + max_seqlen_pad, + ) + mla_metadata, num_splits = get_mla_metadata( + forward_batch.seq_lens.to(torch.int32), + self.num_q_heads, + 1, + ) + self.forward_metadata = FlashMLADecodeMetadata( + mla_metadata, + num_splits, + block_kv_indices, + ) + elif forward_batch.forward_mode.is_target_verify(): + seq_lens_cpu = forward_batch.seq_lens_cpu + self.num_draft_tokens + seq_lens = forward_batch.seq_lens + self.num_draft_tokens + + max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE) + block_kv_indices = torch.full( + (bs, max_seqlen_pad), + -1, + dtype=torch.int32, + device=seq_lens.device, + ) + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + seq_lens, + None, + block_kv_indices, + self.req_to_token.stride(0), + max_seqlen_pad, + ) + mla_metadata, num_splits = get_mla_metadata( + seq_lens.to(torch.int32), + self.num_draft_tokens * self.num_q_heads, + 1, + ) + + # Use FlashMLADecodeMetadata which has the attributes forward_extend expects + self.forward_metadata = FlashMLADecodeMetadata( + mla_metadata, + num_splits, + block_kv_indices, + ) + else: + super().init_forward_metadata(forward_batch) + + def init_cuda_graph_state( + self, + max_bs: int, + max_num_tokens: int, + block_kv_indices: Optional[torch.Tensor] = None, + ): + if block_kv_indices is None: + cuda_graph_kv_indices = torch.full( + (max_bs, (self.max_context_len + PAGE_SIZE) // PAGE_SIZE), + 1, + dtype=torch.int32, + device="cuda", + ) + else: + cuda_graph_kv_indices = block_kv_indices + + if self.num_draft_tokens: + self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata( + torch.ones( + max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device + ), + self.num_draft_tokens * self.num_q_heads, + 1, + ) + else: + self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata( + torch.ones( + max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device + ), + self.num_q_heads, + 1, + ) + self.cuda_graph_kv_indices = cuda_graph_kv_indices + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + ): + if forward_mode.is_decode_or_idle(): + max_seqlen_pad = triton.cdiv(seq_lens.max().item(), PAGE_SIZE) + + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + None, + self.cuda_graph_kv_indices, + self.req_to_token.stride(0), + self.cuda_graph_kv_indices.stride(0), + ) + mla_metadata, num_splits = get_mla_metadata( + seq_lens.to(torch.int32), + self.num_q_heads, + 1, + ) + self.cuda_graph_mla_metadata.copy_(mla_metadata) + self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) + self.forward_metadata = FlashMLADecodeMetadata( + self.cuda_graph_mla_metadata, + self.cuda_graph_num_splits[: bs + 1], + self.cuda_graph_kv_indices[:bs, :max_seqlen_pad], + ) + elif forward_mode.is_target_verify(): + seq_lens = seq_lens + self.num_draft_tokens + max_seqlen_pad = triton.cdiv(seq_lens.max().item(), PAGE_SIZE) + + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + None, + self.cuda_graph_kv_indices, + self.req_to_token.stride(0), + self.cuda_graph_kv_indices.stride(0), + ) + mla_metadata, num_splits = get_mla_metadata( + seq_lens.to(torch.int32), + self.num_draft_tokens * self.num_q_heads, + 1, + ) + self.cuda_graph_mla_metadata.copy_(mla_metadata) + self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) + self.forward_metadata = FlashMLADecodeMetadata( + self.cuda_graph_mla_metadata, + self.cuda_graph_num_splits[: bs + 1], + self.cuda_graph_kv_indices[:bs, :max_seqlen_pad], + ) + else: + super().init_forward_metadata_capture_cuda_graph( + bs, + num_tokens, + req_pool_indices, + seq_lens, + encoder_lens, + forward_mode, + spec_info, + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + seq_lens_cpu: Optional[torch.Tensor], + ): + + if forward_mode.is_decode_or_idle(): + assert seq_lens_cpu is not None + seq_lens = seq_lens[:bs] + seq_lens_cpu = seq_lens_cpu[:bs] + max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE) + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices[:bs], + seq_lens, + None, + self.cuda_graph_kv_indices, + self.req_to_token.stride(0), + self.cuda_graph_kv_indices.stride(0), + ) + mla_metadata, num_splits = get_mla_metadata( + seq_lens.to(torch.int32), + self.num_q_heads, + 1, + ) + self.cuda_graph_mla_metadata.copy_(mla_metadata) + self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) + self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata + self.forward_metadata.num_splits = self.cuda_graph_num_splits[: bs + 1] + self.forward_metadata.block_kv_indices = self.cuda_graph_kv_indices[ + :bs, :max_seqlen_pad + ] + elif forward_mode.is_target_verify(): + seq_lens = seq_lens[:bs] + self.num_draft_tokens + seq_lens_cpu = seq_lens_cpu[:bs] + self.num_draft_tokens + max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE) + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices[:bs], + seq_lens, + None, + self.cuda_graph_kv_indices, + self.req_to_token.stride(0), + self.cuda_graph_kv_indices.stride(0), + ) + mla_metadata, num_splits = get_mla_metadata( + seq_lens.to(torch.int32), + self.num_draft_tokens * self.num_q_heads, + 1, + ) + self.cuda_graph_mla_metadata.copy_(mla_metadata) + self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) + self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata + self.forward_metadata.num_splits = self.cuda_graph_num_splits[: bs + 1] + self.forward_metadata.block_kv_indices = self.cuda_graph_kv_indices[ + :bs, :max_seqlen_pad + ] + else: + super().init_forward_metadata_replay_cuda_graph( + bs, + req_pool_indices, + seq_lens, + seq_lens_sum, + encoder_lens, + forward_mode, + spec_info, + seq_lens_cpu, + ) + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + ): + cache_loc = forward_batch.out_cache_loc + + if k is not None: + assert v is not None + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, + cache_loc, + k, + v, + ) + bs = forward_batch.batch_size + k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + + reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) + if self.data_type == torch.float8_e4m3fn: + reshape_q_fp8 = reshape_q.to(torch.float8_e4m3fn) + o, _ = flash_mla_with_kvcache( + q=reshape_q_fp8, + k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), + block_table=self.forward_metadata.block_kv_indices[:bs], + cache_seqlens=forward_batch.seq_lens.to(torch.int32), + head_dim_v=self.kv_lora_rank, # TODO Retrieve from config. + tile_scheduler_metadata=self.forward_metadata.flashmla_metadata, + num_splits=self.forward_metadata.num_splits, + softmax_scale=layer.scaling, + causal=True, + descale_q=torch.ones((1), dtype=torch.float32, device=reshape_q.device), + descale_k=torch.ones((1), dtype=torch.float32, device=reshape_q.device), + ) + + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + else: + # todo: need check all causal True or False? + o, _ = flash_mla_with_kvcache( + q=reshape_q, + k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), + block_table=self.forward_metadata.block_kv_indices[:bs], + cache_seqlens=forward_batch.seq_lens.to(torch.int32), + head_dim_v=self.kv_lora_rank, # TODO Retrieve from config. + tile_scheduler_metadata=self.forward_metadata.flashmla_metadata, + num_splits=self.forward_metadata.num_splits, + softmax_scale=layer.scaling, + causal=True, + ) + + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + ): + if ( + forward_batch.forward_mode == ForwardMode.EXTEND + or forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND + ): + return super().forward_extend(q, k, v, layer, forward_batch, save_kv_cache) + else: + cache_loc = forward_batch.out_cache_loc + + if k is not None: + assert v is not None + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) + + bs = forward_batch.batch_size + k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + + reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) + if self.data_type == torch.float8_e4m3fn: + reshape_q_fp8 = reshape_q.to(torch.float8_e4m3fn) + o, _ = flash_mla_with_kvcache( + q=reshape_q_fp8, + k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), + block_table=self.forward_metadata.block_kv_indices[:bs], + cache_seqlens=forward_batch.seq_lens.to(torch.int32) + + self.num_draft_tokens, + head_dim_v=self.kv_lora_rank, + tile_scheduler_metadata=self.forward_metadata.flashmla_metadata, + num_splits=self.forward_metadata.num_splits, + softmax_scale=layer.scaling, + causal=True, + descale_q=torch.ones( + (1), dtype=torch.float32, device=reshape_q.device + ), + descale_k=torch.ones( + (1), dtype=torch.float32, device=reshape_q.device + ), + ) + else: + o, _ = flash_mla_with_kvcache( + q=reshape_q, + k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), + block_table=self.forward_metadata.block_kv_indices[:bs], + cache_seqlens=forward_batch.seq_lens.to(torch.int32) + + self.num_draft_tokens, + head_dim_v=self.kv_lora_rank, + tile_scheduler_metadata=self.forward_metadata.flashmla_metadata, + num_splits=self.forward_metadata.num_splits, + softmax_scale=layer.scaling, + causal=True, + ) + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + +# TODO: multi step kv indices optimization +class FlashMLAMultiStepDraftBackend: + """ + Wrap multiple flashmla attention backends as one for multiple consecutive + draft decoding steps. + """ + + def __init__( + self, + model_runner: ModelRunner, + topk: int, + speculative_num_steps: int, + ): + if topk > 1: + raise ValueError( + "Currently FlashMLA only supports topk=1 for speculative decoding" + ) + self.topk = topk + self.speculative_num_steps = speculative_num_steps + max_bs = model_runner.req_to_token_pool.size * self.topk + self.kv_indptr = torch.zeros( + ( + self.speculative_num_steps, + max_bs + 1, + ), + dtype=torch.int32, + device=model_runner.device, + ) + + self.attn_backends = [] + for i in range(self.speculative_num_steps): + self.attn_backends.append( + FlashMLABackend( + model_runner, + skip_prefill=True, + kv_indptr_buf=self.kv_indptr[i], + kv_last_page_len_buf=None, + ) + ) + + def common_template( + self, + forward_batch: ForwardBatch, + call_fn: Callable, + ): + assert forward_batch.spec_info is not None + + for i in range(self.speculative_num_steps - 1): + call_fn(i, forward_batch) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + def call_fn(i, forward_batch): + assert forward_batch.spec_info is not None + self.attn_backends[i].init_forward_metadata(forward_batch) + + self.common_template(forward_batch, call_fn) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + for i in range(self.speculative_num_steps): + self.attn_backends[i].init_cuda_graph_state( + max_bs, max_num_tokens, block_kv_indices=None + ) + + def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_capture_cuda_graph( + forward_batch.batch_size, + forward_batch.batch_size * self.topk, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + ) + + self.common_template(forward_batch, call_fn) + + def init_forward_metadata_replay_cuda_graph( + self, forward_batch: ForwardBatch, bs: int + ): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_replay_cuda_graph( + bs, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + seq_lens_sum=-1, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + seq_lens_cpu=forward_batch.seq_lens_cpu, + ) + + self.common_template(forward_batch, call_fn) diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py new file mode 100644 index 000000000..580a977ec --- /dev/null +++ b/python/sglang/srt/layers/attention/hybrid_attn_backend.py @@ -0,0 +1,139 @@ +from typing import Optional, Union + +import torch + +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput + + +class HybridAttnBackend(AttentionBackend): + """Support different backends for prefill and decode.""" + + def __init__( + self, + model_runner: ModelRunner, + prefill_backend: AttentionBackend, + decode_backend: AttentionBackend, + ): + self.model_runner = model_runner + self.prefill_backend = prefill_backend + self.decode_backend = decode_backend + + def _select_backend(self, forward_mode: ForwardMode) -> AttentionBackend: + """ + Select the appropriate attention backend based on the forward mode. + + Args: + forward_mode: The current forward mode indicating the operation type + + Returns: + The selected attention backend (prefill or decode) + + Note: + - decode_or_idle: Always uses decode backend + - target_verify or draft_extend: Uses decode backend if speculative_attention_mode is "decode", otherwise prefill backend + - prefill: Always uses prefill backend + """ + if forward_mode.is_decode_or_idle(): + return self.decode_backend + elif forward_mode.is_target_verify() or forward_mode.is_draft_extend(): + return ( + self.decode_backend + if self.model_runner.server_args.speculative_attention_mode == "decode" + else self.prefill_backend + ) + else: + return self.prefill_backend + + def init_forward_metadata(self, forward_batch: ForwardBatch): + backend = self._select_backend(forward_batch.forward_mode) + backend.init_forward_metadata(forward_batch) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.decode_backend.init_cuda_graph_state(max_bs, max_num_tokens) + if ( + self.model_runner.server_args.speculative_algorithm is not None + and self.model_runner.server_args.speculative_attention_mode == "prefill" + ): + # When speculative decoding is enabled, we need to initialize the backend + # that will be used for target_verify. + self.prefill_backend.init_cuda_graph_state(max_bs, max_num_tokens) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + backend = self._select_backend(forward_mode) + backend.init_forward_metadata_capture_cuda_graph( + bs, + num_tokens, + req_pool_indices, + seq_lens, + encoder_lens, + forward_mode, + spec_info, + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + ): + backend = self._select_backend(forward_mode) + backend.init_forward_metadata_replay_cuda_graph( + bs, + req_pool_indices, + seq_lens, + seq_lens_sum, + encoder_lens, + forward_mode, + spec_info, + seq_lens_cpu, + ) + + def get_cuda_graph_seq_len_fill_value(self): + return self.decode_backend.get_cuda_graph_seq_len_fill_value() + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + return self.decode_backend.forward_decode( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + backend = self._select_backend(forward_batch.forward_mode) + return backend.forward_extend( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py new file mode 100644 index 000000000..a676573f2 --- /dev/null +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -0,0 +1,584 @@ +from dataclasses import astuple, dataclass +from functools import lru_cache +from typing import Optional, Union + +import torch +import torch.nn.functional as F + +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule +from sglang.srt.layers.attention.fla.fused_recurrent import ( + fused_recurrent_gated_delta_rule_update, +) +from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( + fused_sigmoid_gating_delta_rule_update, +) +from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.models.qwen3_next import Qwen3HybridLinearDecoderLayer, fused_gdn_gating +from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput + + +@dataclass +class ForwardMetadata: + query_start_loc: Optional[torch.Tensor] + mamba_cache_indices: torch.Tensor + + +class MambaAttnBackend(AttentionBackend): + """Attention backend using Mamba kernel.""" + + def __init__(self, model_runner: ModelRunner): + super().__init__() + self.pad_slot_id = -1 # Default pad slot id + self.device = model_runner.device + self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool + self.forward_metadata: ForwardMetadata = None + self.state_indices_list = [] + self.query_start_loc_list = [] + + @classmethod + @lru_cache(maxsize=128) + def _get_cached_arange(cls, bs: int, device_str: str) -> torch.Tensor: + """Cache torch.arange tensors for common batch sizes to avoid repeated allocation.""" + device = torch.device(device_str) + return torch.arange(0, bs + 1, dtype=torch.int32, device=device) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + bs = forward_batch.batch_size + if forward_batch.forward_mode.is_decode_or_idle(): + query_start_loc = self._get_cached_arange(bs, str(self.device)) + elif forward_batch.forward_mode.is_extend(): + if forward_batch.forward_mode.is_target_verify(): + query_start_loc = torch.arange( + 0, + forward_batch.input_ids.shape[0] + 1, + step=forward_batch.spec_info.draft_token_num, + dtype=torch.int32, + device=forward_batch.input_ids.device, + ) + else: + query_start_loc = torch.empty( + (bs + 1,), dtype=torch.int32, device=self.device + ) + query_start_loc[:bs] = forward_batch.extend_start_loc + query_start_loc[bs] = ( + forward_batch.extend_start_loc[-1] + + forward_batch.extend_seq_lens[-1] + ) + else: + raise ValueError(f"Invalid forward mode: {forward_batch.forward_mode=}") + mamba_cache_indices = self.req_to_token_pool.get_mamba_indices( + forward_batch.req_pool_indices + ) + self.forward_metadata = ForwardMetadata( + query_start_loc=query_start_loc, + mamba_cache_indices=mamba_cache_indices, + ) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + for i in range(max_bs): + self.state_indices_list.append( + torch.full((i + 1,), self.pad_slot_id, dtype=torch.int32, device="cuda") + ) + self.query_start_loc_list.append( + torch.empty((i + 2,), dtype=torch.int32, device="cuda") + ) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + if forward_mode.is_decode_or_idle(): + self.query_start_loc_list[bs - 1].copy_(self._get_cached_arange(bs, "cuda")) + elif forward_mode.is_target_verify(): + self.query_start_loc_list[bs - 1].copy_( + torch.arange( + 0, + bs * spec_info.draft_token_num + 1, + step=spec_info.draft_token_num, + dtype=torch.int32, + device="cuda", + ) + ) + else: + raise ValueError(f"Invalid forward mode: {forward_mode=}") + mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) + self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) + self.forward_metadata = ForwardMetadata( + query_start_loc=self.query_start_loc_list[bs - 1], + mamba_cache_indices=self.state_indices_list[bs - 1], + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + ): + num_padding = torch.count_nonzero( + seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value() + ) + # Make sure forward metadata is correctly handled for padding reqs + req_pool_indices[bs - num_padding :] = 0 + mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) + mamba_indices[bs - num_padding :] = -1 + self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) + if forward_mode.is_decode_or_idle(): + self.query_start_loc_list[bs - 1].copy_(self._get_cached_arange(bs, "cuda")) + if num_padding > 0: + self.query_start_loc_list[bs - 1][bs - num_padding :] = bs - num_padding + elif forward_mode.is_target_verify(): + self.query_start_loc_list[bs - 1].copy_( + torch.arange( + 0, + bs * spec_info.draft_token_num + 1, + step=spec_info.draft_token_num, + dtype=torch.int32, + device="cuda", + ) + ) + if num_padding > 0: + self.query_start_loc_list[bs - 1][bs - num_padding :] = ( + bs - num_padding + ) * spec_info.draft_token_num + else: + raise ValueError(f"Invalid forward mode: {forward_mode=}") + + self.forward_metadata = ForwardMetadata( + query_start_loc=self.query_start_loc_list[bs - 1], + mamba_cache_indices=self.state_indices_list[bs - 1], + ) + + def get_cuda_graph_seq_len_fill_value(self): + return 1 # Mamba attn does not use seq lens to index kv cache + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + mixed_qkv = kwargs["mixed_qkv"] + conv_weights = kwargs["conv_weights"] + bias = kwargs["bias"] + activation = kwargs["activation"] + key_dim = kwargs["key_dim"] + value_dim = kwargs["value_dim"] + attn_tp_size = kwargs["attention_tp_size"] + head_k_dim = kwargs["head_k_dim"] + head_v_dim = kwargs["head_v_dim"] + a = kwargs["a"] + b = kwargs["b"] + A_log = kwargs["A_log"] + dt_bias = kwargs["dt_bias"] + layer_id = kwargs["layer_id"] + + conv_states, ssm_states, *rest = self.req_to_token_pool.get_mamba_params( + layer_id + ) + query_start_loc = self.forward_metadata.query_start_loc + cache_indices = self.forward_metadata.mamba_cache_indices + + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_states, + conv_weights, + bias, + activation, + conv_state_indices=cache_indices, + ) + + query, key, value = torch.split( + mixed_qkv, + [ + key_dim // attn_tp_size, + key_dim // attn_tp_size, + value_dim // attn_tp_size, + ], + dim=-1, + ) + # Reshape from [l, h*d] to [1, l, h, d] + seq_len = query.shape[0] + num_heads = query.shape[1] // head_k_dim + query = query.view(1, seq_len, num_heads, head_k_dim) + key = key.view(1, seq_len, num_heads, head_k_dim) + value = value.view(1, seq_len, value.shape[1] // head_v_dim, head_v_dim) + + core_attn_out = fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + dt_bias=dt_bias, + q=query, + k=key, + v=value, + a=a, + b=b, + initial_state_source=ssm_states, + initial_state_indices=cache_indices, + cu_seqlens=query_start_loc, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + ) + + return core_attn_out + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + mixed_qkv = kwargs["mixed_qkv"] + conv_weights = kwargs["conv_weights"] + bias = kwargs["bias"] + activation = kwargs["activation"] + key_dim = kwargs["key_dim"] + value_dim = kwargs["value_dim"] + attn_tp_size = kwargs["attention_tp_size"] + head_k_dim = kwargs["head_k_dim"] + head_v_dim = kwargs["head_v_dim"] + a = kwargs["a"] + b = kwargs["b"] + A_log = kwargs["A_log"] + dt_bias = kwargs["dt_bias"] + layer_id = kwargs["layer_id"] + seq_len = kwargs["seq_len"] + + is_target_verify = forward_batch.forward_mode.is_target_verify() + + query_start_loc = self.forward_metadata.query_start_loc + cache_indices = self.forward_metadata.mamba_cache_indices + + if is_target_verify: + ( + conv_states, + ssm_states, + intermediate_state_cache, + intermediate_conv_window_cache, + ) = self.req_to_token_pool.get_mamba_params(layer_id) + has_initial_states = torch.ones( + seq_len // forward_batch.spec_info.draft_token_num, + dtype=torch.bool, + device=forward_batch.input_ids.device, + ) + conv_states_to_use = conv_states.clone() + else: + conv_states, ssm_states, *rest = self.req_to_token_pool.get_mamba_params( + layer_id + ) + has_initial_states = forward_batch.extend_prefix_lens > 0 + conv_states_to_use = conv_states + + if is_target_verify: + batch_size = seq_len // forward_batch.spec_info.draft_token_num + draft_token_num = forward_batch.spec_info.draft_token_num + mixed_qkv_reshaped = ( + mixed_qkv.view(batch_size, draft_token_num, -1) + .transpose(1, 2) + .contiguous() + ) + mixed_qkv_processed = causal_conv1d_update( + mixed_qkv_reshaped, + conv_states_to_use, + conv_weights, + bias, + activation, + conv_state_indices=cache_indices[:batch_size], + intermediate_conv_window=intermediate_conv_window_cache, + ) + mixed_qkv = ( + mixed_qkv_processed.transpose(1, 2).contiguous().view(seq_len, -1) + ) + else: + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + conv_weights, + bias, + activation=activation, + conv_states=conv_states_to_use, + has_initial_state=has_initial_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + ).transpose(0, 1)[:seq_len] + + key_split_dim = key_dim // attn_tp_size + value_split_dim = value_dim // attn_tp_size + + query, key, value = torch.split( + mixed_qkv, + [key_split_dim, key_split_dim, value_split_dim], + dim=-1, + ) + + actual_seq_len = query.shape[0] + num_heads = query.shape[1] // head_k_dim + num_value_heads = value.shape[1] // head_v_dim + + query = query.view(1, actual_seq_len, num_heads, head_k_dim) + key = key.view(1, actual_seq_len, num_heads, head_k_dim) + value = value.view(1, actual_seq_len, num_value_heads, head_v_dim) + + beta = b.sigmoid() + g = fused_gdn_gating(A_log, a, dt_bias) + + g = g.unsqueeze(0) + beta = beta.unsqueeze(0) + + if is_target_verify: + core_attn_out = fused_recurrent_gated_delta_rule_update( + q=query, + k=key, + v=value, + g=g, + beta=beta, + initial_state_source=ssm_states, + initial_state_indices=cache_indices, + cu_seqlens=query_start_loc, + use_qk_l2norm_in_kernel=True, + disable_state_update=True, + intermediate_states_buffer=intermediate_state_cache, + cache_steps=forward_batch.spec_info.draft_token_num, + ) + else: + recurrent_state = ssm_states[cache_indices] + core_attn_out, last_recurrent_state = chunk_gated_delta_rule( + q=query, + k=key, + v=value, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=True, + cu_seqlens=query_start_loc, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False) + ssm_states[cache_indices] = last_recurrent_state + + return core_attn_out + + +class HybridLinearAttnBackend(AttentionBackend): + """Support different backends for prefill and decode.""" + + def __init__( + self, + full_attn_backend: AttentionBackend, + linear_attn_backend: AttentionBackend, + full_attn_layers: list[int], + ): + self.full_attn_layers = full_attn_layers + self.attn_backend_list = [full_attn_backend, linear_attn_backend] + + def init_forward_metadata(self, forward_batch: ForwardBatch): + for attn_backend in self.attn_backend_list: + attn_backend.init_forward_metadata(forward_batch) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + for attn_backend in self.attn_backend_list: + attn_backend.init_cuda_graph_state(max_bs, max_num_tokens) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + for attn_backend in self.attn_backend_list: + attn_backend.init_forward_metadata_capture_cuda_graph( + bs, + num_tokens, + req_pool_indices, + seq_lens, + encoder_lens, + forward_mode, + spec_info, + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + ): + for attn_backend in self.attn_backend_list: + attn_backend.init_forward_metadata_replay_cuda_graph( + bs, + req_pool_indices, + seq_lens, + seq_lens_sum, + encoder_lens, + forward_mode, + spec_info, + seq_lens_cpu, + ) + + def get_cuda_graph_seq_len_fill_value(self): + return self.attn_backend_list[0].get_cuda_graph_seq_len_fill_value() + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + layer_id = layer.layer_id if layer else kwargs["layer_id"] + if layer_id in self.full_attn_layers: + return self.attn_backend_list[0].forward_decode( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + return self.attn_backend_list[1].forward_decode( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + layer_id = layer.layer_id if layer else kwargs["layer_id"] + if layer_id in self.full_attn_layers: + return self.attn_backend_list[0].forward_extend( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + return self.attn_backend_list[1].forward_extend( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + """Run forward on an attention layer.""" + if forward_batch.forward_mode.is_idle(): + if layer is None: + return torch.empty_like(kwargs["z"]) + return q.new_empty(q.shape[0], layer.tp_q_head_num * layer.v_head_dim) + elif forward_batch.forward_mode.is_decode(): + return self.forward_decode( + q, + k, + v, + layer, + forward_batch, + save_kv_cache=save_kv_cache, + **kwargs, + ) + else: + return self.forward_extend( + q, + k, + v, + layer, + forward_batch, + save_kv_cache=save_kv_cache, + **kwargs, + ) + + def update_mamba_state_after_mtp_verify(self, accepted_length, model): + request_number = accepted_length.shape[0] + + state_indices_tensor = self.attn_backend_list[ + 1 + ].forward_metadata.mamba_cache_indices[:request_number] + + mamba_caches = self.attn_backend_list[ + 1 + ].req_to_token_pool.get_mamba_params_all_layers() + + ( + conv_states, + ssm_states, + intermediate_state_cache, + intermediate_conv_window_cache, + ) = mamba_caches + + # SSM state updates (chunked to reduce peak memory) + valid_mask = accepted_length > 0 + + # Compute common indices once to avoid duplication + last_steps_all = (accepted_length - 1).to(torch.int64) + valid_state_indices = state_indices_tensor[valid_mask].to(torch.int64) + last_steps = last_steps_all[valid_mask].to(torch.int64) + + if valid_state_indices.numel() > 0: + chunk = 256 + num_valid = valid_state_indices.numel() + + # SSM state updates + for i in range(0, num_valid, chunk): + idx = valid_state_indices[i : i + chunk] + steps = last_steps[i : i + chunk] + # per (cache line, step) + for j in range(idx.numel()): + ci = idx[j].item() + st = steps[j].item() + ssm_states[:, ci, :].copy_( + intermediate_state_cache[:, ci, st].to( + ssm_states.dtype, copy=False + ) + ) + + # Conv window updates + for i in range(0, num_valid, chunk): + idx = valid_state_indices[i : i + chunk] + steps = last_steps[i : i + chunk] + for j in range(idx.numel()): + ci = idx[j].item() + st = steps[j].item() + conv_states[:, ci, :, :].copy_( + intermediate_conv_window_cache[:, ci, st].to( + conv_states.dtype, copy=False + ) + ) diff --git a/python/sglang/srt/layers/attention/intel_amx_backend.py b/python/sglang/srt/layers/attention/intel_amx_backend.py new file mode 100644 index 000000000..39e5c7428 --- /dev/null +++ b/python/sglang/srt/layers/attention/intel_amx_backend.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + + +class IntelAMXAttnBackend(AttentionBackend): + def __init__(self, model_runner: ModelRunner): + import sgl_kernel + + super().__init__() + self.forward_metadata = None + self.device = model_runner.device + + self.num_head = ( + model_runner.model_config.num_attention_heads // model_runner.tp_size + ) + + self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1] + + self.decode_attention_fwd = torch.ops.sgl_kernel.decode_attention_cpu + self.extend_attention_fwd = torch.ops.sgl_kernel.extend_attention_cpu + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init the metadata for a forward pass.""" + + bs = forward_batch.batch_size + attn_logits = torch.zeros( + ( + bs, + self.num_head, + 8, # self.num_kv_splits, + self.v_head_dim + 1, + ), + dtype=torch.float32, + device=self.device, + ) + if forward_batch.forward_mode.is_decode_or_idle(): + max_extend_len = None + else: + max_extend_len = torch.max(forward_batch.extend_seq_lens).item() + self.forward_metadata = (attn_logits, max_extend_len) + + def get_graph_seq_len_fill_value(self): + return 1 + + def forward_extend( + self, + q, + k, + v, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + + _, max_extend_len = self.forward_metadata + + self.extend_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + k, + v, + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + forward_batch.req_to_token_pool.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.extend_seq_lens, + forward_batch.extend_start_loc, + max_extend_len, + layer.scaling, + layer.logit_cap, + ) + return o + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + attn_logits, _ = self.forward_metadata + + q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim) + + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + self.decode_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + k, + v, + forward_batch.out_cache_loc, + attn_logits, + forward_batch.req_to_token_pool.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + layer.scaling, + layer.logit_cap, + ) + + return o + + def support_triton(self): + return False diff --git a/python/sglang/srt/layers/attention/mamba/causal_conv1d.py b/python/sglang/srt/layers/attention/mamba/causal_conv1d.py new file mode 100644 index 000000000..d004337ff --- /dev/null +++ b/python/sglang/srt/layers/attention/mamba/causal_conv1d.py @@ -0,0 +1,128 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2024, Tri Dao. +# Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py + +from typing import Optional + +import torch +from sgl_kernel import causal_conv1d_fwd +from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel + +PAD_SLOT_ID = -1 + + +def causal_conv1d_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + query_start_loc: Optional[torch.Tensor] = None, + cache_indices: Optional[torch.Tensor] = None, + has_initial_state: Optional[torch.Tensor] = None, + conv_states: Optional[torch.Tensor] = None, + activation: Optional[str] = "silu", + pad_slot_id: int = PAD_SLOT_ID, +): + """ + x: (batch, dim, seqlen) or (dim,cu_seq_len) for varlen + sequences are concatenated from left to right for varlen + weight: (dim, width) + bias: (dim,) + query_start_loc: (batch + 1) int32 + The cumulative sequence lengths of the sequences in + the batch, used to index into sequence. prepended by 0. + for example: query_start_loc = torch.Tensor([0,10,16,17]), + x.shape=(dim,17) + cache_indices: (batch) int32 + indicates the corresponding state index, + like so: conv_state = conv_states[cache_indices[batch_id]] + has_initial_state: (batch) bool + indicates whether should the kernel take the current state as initial + state for the calculations + conv_states: (...,dim,width - 1) itype + updated inplace if provided + activation: either None or "silu" or "swish" + pad_slot_id: int + if cache_indices is passed, lets the kernel identify padded + entries that will not be processed, + for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] + in this case, the kernel will not process entries at + indices 0 and 3 + + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + if x.stride(-1) != 1: + x = x.contiguous() + bias = bias.contiguous() if bias is not None else None + + causal_conv1d_fwd( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + activation in ["silu", "swish"], + pad_slot_id, + ) + return x + + +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + activation: Optional[str] = None, + cache_seqlens: Optional[torch.Tensor] = None, + conv_state_indices: Optional[torch.Tensor] = None, + pad_slot_id: int = PAD_SLOT_ID, +): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state + starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If not None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + pad_slot_id: int + if cache_indices is passed, lets the kernel identify padded + entries that will not be processed, + for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id] + in this case, the kernel will not process entries at + indices 0 and 3 + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError( + f"activation must be None, silu, or swish, actual: {activation}" + ) + activation_val = activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + causal_conv1d_update_kernel( + x, + conv_state, + weight, + bias, + activation_val, + cache_seqlens, + conv_state_indices, + pad_slot_id, + ) + if unsqueeze: + x = x.squeeze(-1) + return x diff --git a/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py b/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py new file mode 100644 index 000000000..3c1bdec48 --- /dev/null +++ b/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py @@ -0,0 +1,1052 @@ +# Copyright (c) 2024, Tri Dao. +# Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py +# and https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/ops/causal_conv1d.py + +from typing import Optional, Union + +import numpy as np +import torch + +PAD_SLOT_ID = -1 +import triton +import triton.language as tl + + +@triton.jit() +def _causal_conv1d_fwd_kernel( # continuous batching + # Pointers to matrices + x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences + w_ptr, # (dim, width) + bias_ptr, + initial_states_ptr, # conv_states_ptr + cache_indices_ptr, # conv_state_indices_ptr + has_initial_states_ptr, + query_start_loc_ptr, + batch_ptr, + token_chunk_offset_ptr, + o_ptr, # (dim, seqlen) - actually pointing to x_ptr + # Matrix dimensions + batch: tl.int32, # actually padded_batch + dim: tl.constexpr, + seqlen: tl.int32, # cu_seqlen + num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines + # Strides + stride_x_seq: tl.constexpr, # stride to get to next sequence, + stride_x_dim: tl.constexpr, # stride to get to next feature-value, + stride_x_token: tl.constexpr, # stride to get to next token (same feature-index, same sequence-index) + stride_w_dim: tl.constexpr, # stride to get to next dim-axis value + stride_w_width: tl.constexpr, # stride to get to next width-axis value + stride_istate_seq: tl.constexpr, + stride_istate_dim: tl.constexpr, + stride_istate_token: tl.constexpr, + stride_o_seq: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.constexpr, + # others + pad_slot_id: tl.constexpr, + # Meta-parameters + HAS_BIAS: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + HAS_INITIAL_STATES: tl.constexpr, + HAS_CACHE: tl.constexpr, + IS_CONTINUOUS_BATCHING: tl.constexpr, + USE_PAD_SLOT: tl.constexpr, + NP2_STATELEN: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + conv_states_ptr = initial_states_ptr + conv_state_indices_ptr = cache_indices_ptr + stride_conv_state_seq = stride_istate_seq + stride_conv_state_dim = stride_istate_dim + stride_conv_state_tok = stride_istate_token + state_len = ( + KERNEL_WIDTH - 1 + ) # can be passed via argument if it's not the same as this value + + # one program handles one chunk in a single sequence + # rather than mixing sequences - to make updating initial_states across sequences efficiently + + # single-sequence id + idx_seq = tl.load(batch_ptr + tl.program_id(0)) + chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0)) + + # BLOCK_N elements along the feature-dimension (channel) + idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + + if idx_seq == pad_slot_id: + return + + sequence_start_index = tl.load(query_start_loc_ptr + idx_seq) + sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1) + # find the actual sequence length + seqlen = sequence_end_index - sequence_start_index + + token_offset = BLOCK_M * chunk_offset + segment_len = min(BLOCK_M, seqlen - token_offset) + + # base of the sequence + x_base = ( + x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim + ) # [BLOCK_N,] + + if IS_CONTINUOUS_BATCHING: + # cache_idx + conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq).to(tl.int64) + else: + # cache_idx + conv_state_batch_coord = idx_seq + if USE_PAD_SLOT: # noqa + if conv_state_batch_coord == pad_slot_id: + # not processing as this is not the actual sequence + return + conv_states_base = ( + conv_states_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + ) # [BLOCK_N,] + + w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] + + # Does 2 things: + # 1. READ prior-block init-state data - [done by every Triton programs] + # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0] + if chunk_offset == 0: + # read from conv_states + load_init_state = False + if HAS_INITIAL_STATES: # the new HAS_INITIAL_STATES + load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1) + if load_init_state: + # load from conv_states + prior_tokens = conv_states_base + (state_len - 1) * stride_conv_state_tok + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 5: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + else: + # prior-tokens are zeros + if KERNEL_WIDTH >= 2: # STRATEGY1 + # first chunk and does not have prior-token, so just set to 0 + col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 3: # STRATEGY1 + col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 4: # STRATEGY1 + col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 5: # STRATEGY1 + col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + + # STEP 2: + # here prepare data for updating conv_state + if ( + state_len <= seqlen + ): # SMALL_CACHE=True (only move part of 'x' into conv_state cache) + # just read from 'x' + # copy 'x' data to conv_state + # load only 'x' data (and set 0 before 'x' if seqlen < state_len) + idx_tokens_last = (seqlen - state_len) + tl.arange( + 0, NP2_STATELEN + ) # [BLOCK_M] + x_ptrs = ( + x_ptr + + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None] + + (idx_feats * stride_x_dim)[None, :] + ) # [BLOCK_M,BLOCK_N,] + mask_x = ( + (idx_tokens_last >= 0)[:, None] + & (idx_tokens_last < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + new_conv_state = tl.load(x_ptrs, mask_x, 0.0) + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + conv_states_ptrs_target = ( + conv_states_base[None, :] + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) + + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.debug_barrier() # NOTE: use this due to bug in Triton compiler + tl.store(conv_states_ptrs_target, new_conv_state, mask) + + else: + if load_init_state: + # update conv_state by shifting left, i.e. take last few cols from conv_state + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + conv_states_ptrs_source = ( + conv_states_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = ( + (conv_state_batch_coord < num_cache_lines) + & ((idx_tokens_conv + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0) + + VAL = state_len - seqlen + + x_ptrs = ( + x_base[None, :] + + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + + tl.debug_barrier() # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load + new_conv_state = tl.where( + mask, conv_state, loaded_x + ) # BUG in 'tl.where' which requires a barrier before this + conv_states_ptrs_target = ( + conv_states_base + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ + None, : + ] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + else: # load_init_state == False + # update conv_state by shifting left, BUT + # set cols prior to 'x' as zeros + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + VAL = state_len - seqlen + + x_ptrs = ( + x_base[None, :] + + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + new_conv_state = tl.load(x_ptrs, mask_x, 0.0) + + conv_states_ptrs_target = ( + conv_states_base + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ + None, : + ] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + + else: # chunk_offset > 0 + # read prior-token data from `x` + load_init_state = True + prior_tokens = x_base + (token_offset - 1) * stride_x_token + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 5: + # ruff: noqa: F841 + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + + if HAS_BIAS: + bias = bias_ptr + idx_feats + mask_bias = idx_feats < dim + acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to( + tl.float32 + ) # [BLOCK_N] + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + + x_base_1d = x_base + token_offset * stride_x_token # starting of chunk + + # PRE-LOAD WEIGHTS + mask_w = idx_feats < dim + if KERNEL_WIDTH >= 2: + w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor + w_col0 = tl.load(w_ptrs, mask_w, other=0.0) + w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor + w_col1 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 3: + w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor + w_col2 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 4: + w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor + w_col3 = tl.load(w_ptrs, mask_w, other=0.0) + mask_x_1d = idx_feats < dim + for idx_token in range(segment_len): + acc = acc_preload + + matrix_w = w_col0 + matrix_x = col0 + for j in tl.static_range(KERNEL_WIDTH): + + if KERNEL_WIDTH == 2: + if j == 1: # KERNEL_WIDTH-1: + matrix_w = w_col1 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 3: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 4: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + + acc += matrix_x * matrix_w # [BLOCK_N] + + if KERNEL_WIDTH == 2: + col0 = matrix_x + elif KERNEL_WIDTH == 3: + col0 = col1 + col1 = matrix_x + elif KERNEL_WIDTH == 4: + col0 = col1 + col1 = col2 + col2 = matrix_x + + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + mask_1d = (idx_token < segment_len) & ( + idx_feats < dim + ) # token-index # feature-index + o_ptrs = ( + o_ptr + + (sequence_start_index + token_offset + idx_token) * stride_o_token + + (idx_feats * stride_o_dim) + ) + + tl.store(o_ptrs, acc, mask=mask_1d) + + +def causal_conv1d_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: Union[torch.Tensor, None], + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: Optional[torch.Tensor] = None, + has_initial_state: Optional[torch.Tensor] = None, + activation: Optional[str] = "silu", + pad_slot_id: int = PAD_SLOT_ID, + metadata=None, + validate_data=False, +): + """support varlen + continuous batching when x is 2D tensor + + x: (dim,cu_seq_len) + cu_seq_len = total tokens of all seqs in that batch + sequences are concatenated from left to right for varlen + weight: (dim, width) + conv_states: (...,dim,width - 1) itype + updated inplace if provided + [it use `cache_indices` to get the index to the cache of conv_state for that sequence + + conv_state[cache_indices[i]] for seq-i - to be used as initial_state when has_initial_state[i] = True + and after that conv_state[cache_indices[i]] need to be shift-left and updated with values from 'x' + ] + query_start_loc: (batch + 1) int32 + The cumulative sequence lengths of the sequences in + the batch, used to index into sequence. prepended by 0. + if + x = [5, 1, 1, 1] <- continuous batching (batch=4) + then + query_start_loc = [0, 5, 6, 7, 8] <- the starting index of the next sequence; while the last value is + the ending index of the last sequence + [length(query_start_loc)-1 == batch] + for example: query_start_loc = torch.Tensor([0,10,16,17]), + x.shape=(dim,17) + cache_indices: (batch) int32 + indicates the corresponding state index, + like so: conv_state = conv_states[cache_indices[batch_id]] + has_initial_state: (batch) bool + indicates whether should the kernel take the current state as initial + state for the calculations + [single boolean for each sequence in the batch: True or False] + bias: (dim,) + activation: either None or "silu" or "swish" or True + pad_slot_id: int + if cache_indices is passed, lets the kernel identify padded + entries that will not be processed, + for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] + in this case, the kernel will not process entries at + indices 0 and 3 + + out: same shape as `x` + """ + if isinstance(activation, bool) and activation: + activation = "silu" + + args = None + out = torch.empty_like(x) + if metadata is not None: + cu_seqlen = metadata.cu_seqlen + nums_dict = metadata.nums_dict + # x = metadata.x + args = nums_dict + batch_ptr = metadata.batch_ptr + token_chunk_offset_ptr = metadata.token_chunk_offset_ptr + else: + seqlens = np.diff(query_start_loc.to("cpu")) + args = seqlens + MAX_NUM_PROGRAMS = 1024 + + batch_ptr = torch.full( + (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device + ) # tracking which seq-idx the Triton program is handling + token_chunk_offset_ptr = torch.full( + (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device + ) # tracking BLOCK_M-based index in the sequence the Triton program is handling + + is_channel_last = (x.stride(0) == 1) & (x.stride(1) > 1) + dim, cu_seqlen = x.shape + _, width = weight.shape + state_len = width - 1 + np2_statelen = triton.next_power_of_2(state_len) + + padded_batch = query_start_loc.size(0) - 1 + stride_x_seq = 0 + stride_x_dim = x.stride(0) + stride_x_token = x.stride(1) + stride_w_dim = weight.stride(0) + stride_w_width = weight.stride(1) + stride_istate_seq = 0 + stride_istate_dim = 0 + stride_istate_token = 0 + num_cache_lines = 0 + if conv_states is not None: + # extensions to support vLLM: + # 1. conv_states is used to replaced initial_states + # 2. conv_states serve as a cache with num cache lines can be larger than batch size + # 3. mapping from sequence x[idx] to a cache line at index as specified via cache_indices[idx] + # 4. computation can be skipped if cache_indices[idx] == pad_slot_id + num_cache_lines = conv_states.size(0) + assert ( + num_cache_lines == conv_states.shape[0] + and dim == conv_states.shape[1] + and width - 1 <= conv_states.shape[2] + ) + stride_istate_seq = conv_states.stride(0) + stride_istate_dim = conv_states.stride(1) + stride_istate_token = conv_states.stride(2) + # assert stride_istate_dim == 1 + if out.dim() == 2: + stride_o_seq = 0 + stride_o_dim = out.stride(0) + stride_o_token = out.stride(1) + else: + stride_o_seq = out.stride(0) + stride_o_dim = out.stride(1) + stride_o_token = out.stride(2) + + if validate_data: + assert x.dim() == 2 + assert query_start_loc is not None + assert query_start_loc.dim() == 1 + assert x.stride(0) == 1 or x.stride(1) == 1 + if bias is not None: + assert bias.dim() == 1 + assert dim == bias.size(0) + if cache_indices is not None: + assert cache_indices.dim() == 1 + assert padded_batch == cache_indices.size(0) + if has_initial_state is not None: + assert has_initial_state.size() == (padded_batch,) + assert ( + conv_states is not None + ), "ERROR: `has_initial_state` is used, which needs also `conv_states`" + assert weight.stride(1) == 1 + assert (dim, width) == weight.shape + assert is_channel_last, "Need to run in channel-last layout" + + if metadata is None: + + def num_program(META, seqlens): + tot = 0 + + mlist = [] + offsetlist = [] # type: ignore + + nums = -(-seqlens // META["BLOCK_M"]) + + tot = nums.sum().item() + mlist = np.repeat(np.arange(len(nums)), nums) + for idx, num in enumerate(nums): + offsetlist.extend( + range(num) + ) # chunk-idx if a sequence is split into multiple chunks + + if META["batch_ptr"].nelement() < len(mlist): + newlen = len(mlist) + 1 + META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + + if META["batch_ptr"].nelement() >= len(mlist): + META["batch_ptr"][0 : len(mlist)].copy_( + torch.from_numpy(np.array(mlist)) + ) + META["token_chunk_offset_ptr"][0 : len(mlist)].copy_( + torch.from_numpy(np.array(offsetlist)) + ) + + META["batch_ptr"] = META["batch_ptr"].to(META["x_ptr"].device) + META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to( + META["x_ptr"].device + ) + return tot + + else: + + def num_program(META, nums_dict): + tot = nums_dict[META["BLOCK_M"]]["tot"] + + mlist = nums_dict[META["BLOCK_M"]]["mlist"] + mlist_len = nums_dict[META["BLOCK_M"]]["mlist_len"] + + offsetlist = nums_dict[META["BLOCK_M"]]["offsetlist"] + + if nums_dict[META["BLOCK_M"]]["batch_ptr"] is not None: + META["batch_ptr"] = nums_dict[META["BLOCK_M"]]["batch_ptr"] + META["token_chunk_offset_ptr"] = nums_dict[META["BLOCK_M"]][ + "token_chunk_offset_ptr" + ] + else: + if META["batch_ptr"].nelement() < mlist_len: + newlen = mlist_len + 1 + META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + + if META["batch_ptr"].nelement() >= mlist_len: + META["batch_ptr"][0:mlist_len].copy_(mlist) + META["token_chunk_offset_ptr"][0:mlist_len].copy_(offsetlist) + return tot + + def grid(META): + return ( + num_program(META, args), + triton.cdiv(dim, META["BLOCK_N"]), + ) + + if batch_ptr.device != x.device: + batch_ptr = batch_ptr.to(x.device) + token_chunk_offset_ptr = token_chunk_offset_ptr.to(x.device) + + _causal_conv1d_fwd_kernel[grid]( + # Pointers to matrices + x, + weight, + bias, + conv_states, + cache_indices, + has_initial_state, + query_start_loc, + batch_ptr, + token_chunk_offset_ptr, + out, + # Matrix dimensions + padded_batch, + dim, + cu_seqlen, + num_cache_lines, + # stride + stride_x_seq, + stride_x_dim, + stride_x_token, + stride_w_dim, + stride_w_width, + stride_istate_seq, + stride_istate_dim, + stride_istate_token, + stride_o_seq, + stride_o_dim, + stride_o_token, + # others + pad_slot_id, + # META + HAS_BIAS=bias is not None, + KERNEL_WIDTH=width, + SILU_ACTIVATION=activation in ["silu", "swish"], + HAS_INITIAL_STATES=has_initial_state is not None, + HAS_CACHE=conv_states is not None, + IS_CONTINUOUS_BATCHING=cache_indices is not None, + USE_PAD_SLOT=pad_slot_id is not None, + NP2_STATELEN=np2_statelen, + # launch_cooperative_grid=True + BLOCK_M=8, + BLOCK_N=256, + num_stages=2, + ) + return out + + +@triton.jit() +def _causal_conv1d_update_kernel( + # Pointers to matrices + x_ptr, # (batch, dim, seqlen) + w_ptr, # (dim, width) + bias_ptr, + conv_state_ptr, + cache_seqlens_ptr, # circular buffer + conv_state_indices_ptr, + num_accepted_tokens_ptr, + intermediate_conv_window_ptr, + o_ptr, # (batch, dim, seqlen) + # Matrix dimensions + batch: int, + dim: tl.constexpr, + seqlen: tl.constexpr, + state_len: tl.constexpr, + num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines + # Strides + stride_x_seq: tl.constexpr, + stride_x_dim: tl.constexpr, + stride_x_token: tl.constexpr, + stride_w_dim: tl.constexpr, + stride_w_width: tl.constexpr, + stride_conv_state_seq: tl.constexpr, + stride_conv_state_dim: tl.constexpr, + stride_conv_state_tok: tl.constexpr, + stride_state_indices: tl.constexpr, + stride_inter_seq: tl.constexpr, + stride_inter_step: tl.constexpr, + stride_inter_dim: tl.constexpr, + stride_inter_win: tl.constexpr, + stride_o_seq: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.constexpr, + # others + pad_slot_id: tl.constexpr, + # Meta-parameters + HAS_BIAS: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + IS_CONTINUOUS_BATCHING: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + NP2_STATELEN: tl.constexpr, + USE_PAD_SLOT: tl.constexpr, + BLOCK_N: tl.constexpr, + SAVE_INTERMEDIATE: tl.constexpr, +): + # ruff: noqa: E501 + idx_seq = tl.program_id(0) + if idx_seq >= batch: + return + + # [BLOCK_N,] elements along the feature-dimension (channel) + idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + + if IS_CONTINUOUS_BATCHING: + # mask = idx_seq < batch + conv_state_batch_coord = tl.load( + conv_state_indices_ptr + idx_seq * stride_state_indices + ).to(tl.int64) + else: + conv_state_batch_coord = idx_seq + if USE_PAD_SLOT: # noqa + if conv_state_batch_coord == pad_slot_id: + # not processing as this is not the actual sequence + return + + if IS_SPEC_DECODING: + # The rolling of conv state: + # + # Before forward, the conv_state is: + # [history1, history2, ..., historyM]. + # + # After forward, the conv_state becomes: + # [history2, ..., historyM, draft1, draft2, ..., draftN]. + # + # After acceptance, it becomes: + # + # - accept 1 tokens: [history2, ..., historyM, draft1] + # - accept 2 tokens: [history3, ..., historyM, draft1, draft2] + # - and so on. + conv_state_token_offset = tl.load(num_accepted_tokens_ptr + idx_seq) - 1 + else: + conv_state_token_offset = 0 + + # STEP 1: READ init_state data + conv_states_base = ( + conv_state_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + ) + mask_w = idx_feats < dim + + prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok + if KERNEL_WIDTH >= 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 3: + conv_states_ptrs = prior_tokens + 1 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 4: + conv_states_ptrs = prior_tokens + 2 * stride_conv_state_tok # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 5: + conv_states_ptrs = prior_tokens + 3 * stride_conv_state_tok # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + + # STEP 2: assume state_len > seqlen + idx_tokens = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # The conv_state updates works in a sliding window manner, + # at each forward pass, the tokens are shift by 1, so we + # load since idx_tokens + 1. + conv_state_ptrs_source = ( + conv_state_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + conv_state_token_offset * stride_conv_state_tok + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens + 1) * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = ( + (conv_state_batch_coord < num_cache_lines) + & ((idx_tokens + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0) + + VAL = state_len - seqlen + x_base = x_ptr + (idx_seq * stride_x_seq) + (idx_feats * stride_x_dim) # [BLOCK_N] + + x_ptrs = ( + x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens - VAL >= 0)[:, None] + & (idx_tokens - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + tl.debug_barrier() + + new_conv_state = tl.where(mask, conv_state, loaded_x) + + conv_state_base = ( + conv_state_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + ) # [BLOCK_N,] + conv_state_ptrs_target = ( + conv_state_base + (idx_tokens * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.store(conv_state_ptrs_target, new_conv_state, mask) + + # STEP 3: init accumulator + if HAS_BIAS: + bias = bias_ptr + idx_feats + mask_bias = idx_feats < dim + acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to( + tl.float32 + ) # [BLOCK_N] + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + + # STEP 4: + # PRE-LOAD WEIGHTS + # first kernel column, configured for weights to handle BLOCK_N features in range + w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] + mask_w = idx_feats < dim + if KERNEL_WIDTH >= 2: + w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor + w_col0 = tl.load(w_ptrs, mask_w, other=0.0) + w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor + w_col1 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 3: + w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor + w_col2 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 4: + w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor + w_col3 = tl.load(w_ptrs, mask_w, other=0.0) + + x_base_1d = x_base # starting of chunk [BLOCK_N] + mask_x_1d = idx_feats < dim + + # STEP 5: compute each token + for idx_token in tl.static_range(seqlen): + acc = acc_preload + + matrix_w = w_col0 + matrix_x = col0 + for j in tl.static_range(KERNEL_WIDTH): + if KERNEL_WIDTH == 2: + if j == 1: # KERNEL_WIDTH-1: + matrix_w = w_col1 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 3: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 4: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + + acc += matrix_x * matrix_w # [BLOCK_N] + + if KERNEL_WIDTH == 2: + col0 = matrix_x + elif KERNEL_WIDTH == 3: + col0 = col1 + col1 = matrix_x + elif KERNEL_WIDTH == 4: + col0 = col1 + col1 = col2 + col2 = matrix_x + + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + mask_1d = (idx_token < seqlen) & ( + idx_feats < dim + ) # token-index # feature-index + o_ptrs = ( + o_ptr + + (idx_seq) * stride_o_seq + + idx_token * stride_o_token + + (idx_feats * stride_o_dim) + ) + + tl.store(o_ptrs, acc, mask=mask_1d) + + if SAVE_INTERMEDIATE: + # Save the window state after consuming this token + # Layout: [seq(cache line), step, dim, win(K-1)] + base_ptr = ( + intermediate_conv_window_ptr + + conv_state_batch_coord * stride_inter_seq + + idx_token * stride_inter_step + + idx_feats * stride_inter_dim + ) + if KERNEL_WIDTH >= 2: + tl.store(base_ptr + 0 * stride_inter_win, col0, mask=mask_w) + if KERNEL_WIDTH >= 3: + tl.store(base_ptr + 1 * stride_inter_win, col1, mask=mask_w) + if KERNEL_WIDTH >= 4: + tl.store(base_ptr + 2 * stride_inter_win, col2, mask=mask_w) + + +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + activation: Union[bool, str, None] = None, + cache_seqlens: Optional[torch.Tensor] = None, + conv_state_indices: Optional[torch.Tensor] = None, + num_accepted_tokens: Optional[torch.Tensor] = None, + intermediate_conv_window: Optional[torch.Tensor] = None, + pad_slot_id: int = PAD_SLOT_ID, + metadata=None, + validate_data=False, +): + """ + x: (batch, dim) or (batch, dim, seqlen) + [shape=2: single token prediction] + [shape=3: single or multiple tokens prediction] + conv_state: (..., dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state + starting at the index + @cache_seqlens % state_len. + conv_state_indices: (batch,), dtype int32 + If not None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + pad_slot_id: int + if cache_indices is passed, lets the kernel identify padded + entries that will not be processed, + for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id] + in this case, the kernel will not process entries at + indices 0 and 3 + out: (batch, dim) or (batch, dim, seqlen) + """ + if validate_data: + assert cache_seqlens is None # not implemented yet - ok for vLLM + assert pad_slot_id is not None + assert x.stride(1) == 1 + if isinstance(activation, bool): + activation = "silu" if activation is True else None + elif activation is not None: + assert activation in ["silu", "swish"] + unsqueeze = x.dim() == 2 + if unsqueeze: + # make it (batch, dim, seqlen) with seqlen == 1 + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + _, width = weight.shape + # conv_state: (..., dim, state_len), where state_len >= width - 1 + num_cache_lines, _, state_len = conv_state.size() + + if validate_data: + assert dim == weight.size(0) + assert ( + conv_state.stride(-2) == 1 + ), f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})" + assert state_len >= width - 1 + # when above happens, we don't shift-left to keep any records in conv_state + assert dim == conv_state.size(1) + if conv_state_indices is None: + assert conv_state.size(0) >= batch + else: + assert (batch,) == conv_state_indices.shape + + assert num_cache_lines >= batch + assert weight.stride(1) == 1 # Need this + assert cache_seqlens is None # not needed for vLLM - circular buffer + + # adopt the strategy in vLLM that overwrite on 'x' directly, rather than creating a new tensor 'o' + out = x + stride_w_dim, stride_w_width = weight.stride() + + stride_x_seq, stride_x_dim, stride_x_token = x.stride() # X (batch, dim, seqlen) + + stride_o_seq, stride_o_dim, stride_o_token = out.stride() + stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride() + stride_state_indices = ( + conv_state_indices.stride(0) if conv_state_indices is not None else 0 + ) + state_len = width - 1 + (seqlen - 1) # effective state_len needed + np2_statelen = triton.next_power_of_2(state_len) + + def grid(META): + return ( + batch, + triton.cdiv(dim, META["BLOCK_N"]), + ) + + # prepare intermediate buffer strides if provided + if intermediate_conv_window is not None: + stride_inter_seq, stride_inter_step, stride_inter_dim, stride_inter_win = ( + intermediate_conv_window.stride(0), + intermediate_conv_window.stride(1), + intermediate_conv_window.stride(2), + intermediate_conv_window.stride(3), + ) + else: + stride_inter_seq = stride_inter_step = stride_inter_dim = stride_inter_win = 0 + + _causal_conv1d_update_kernel[grid]( + # Pointers to matrices + x, + weight, + bias, + conv_state, + cache_seqlens, + conv_state_indices, + num_accepted_tokens, + intermediate_conv_window if intermediate_conv_window is not None else x, + out, + # Matrix dimensions + batch, + dim, + seqlen, + state_len, + num_cache_lines, + # stride + stride_x_seq, + stride_x_dim, + stride_x_token, + stride_w_dim, + stride_w_width, + stride_istate_seq, + stride_istate_dim, + stride_istate_token, + stride_state_indices, + stride_inter_seq, + stride_inter_step, + stride_inter_dim, + stride_inter_win, + stride_o_seq, + stride_o_dim, + stride_o_token, + # others + pad_slot_id, + # META + HAS_BIAS=bias is not None, + KERNEL_WIDTH=width, + SILU_ACTIVATION=activation in ["silu", "swish"], + IS_CONTINUOUS_BATCHING=conv_state_indices is not None, + IS_SPEC_DECODING=num_accepted_tokens is not None, + NP2_STATELEN=np2_statelen, + USE_PAD_SLOT=pad_slot_id is not None, + BLOCK_N=256, + SAVE_INTERMEDIATE=intermediate_conv_window is not None, + ) + if unsqueeze: + out = out.squeeze(-1) + return out diff --git a/python/sglang/srt/layers/attention/mamba/mamba.py b/python/sglang/srt/layers/attention/mamba/mamba.py new file mode 100644 index 000000000..045a04048 --- /dev/null +++ b/python/sglang/srt/layers/attention/mamba/mamba.py @@ -0,0 +1,64 @@ +from typing import Callable, List, Tuple + +import torch + +LoaderFunction = Callable[[torch.Tensor, torch.Tensor], None] + + +def mamba_v2_sharded_weight_loader( + shard_spec: List[Tuple[int, int, float]], + tp_size: int, + tp_rank: int, +) -> LoaderFunction: + """Create a weight loader for mamba v2. This ensures that the projections + are correctly sharded so that they can be split into x, B, C. It also + ensures the the all the groups corresponding to a head shard is placed + together with it. + """ + + def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + + # - track boundary of (sharded) param, and loaded_weight, respectively + boundary, loaded_boundary = 0, 0 + + # - iterate over the shard specs + for full_dim, extra, duplicate_groups in shard_spec: + # - full dim is the model dim (before TP). + # - extra > 0, means there is expected overall increase + # of dimensions. This is so because of replication. + # - ratio is used map the tp_rank to the actual shard + # rank. This is useful when there is replication of + # groups to accompany head shards. + + # - size of the loaded shard + shard_size = full_dim // tp_size + + # - compute the rank into the loaded shard. + # - if there is replication, different TP shards will + # take from the same rank. + # NOTE: currently we only support duplication + # in the case where num_groups == 1 + rank = 0 if duplicate_groups else tp_rank + + # - leftmost boundary index into loaded weight. + loaded_skip = rank * shard_size + loaded_start_idx = loaded_boundary + loaded_skip + + # - take these many dims from the loaded weight. + take = min(shard_size, full_dim - extra - loaded_skip) + + # - always shard on dim 0 + # - the ignore is for a mundane mypy error as it does not + # seem to handle slices well. + # https://github.com/python/mypy/issues/2410 + param.data[ + boundary : (boundary + take), ... # type: ignore[misc] + ] = loaded_weight[ + loaded_start_idx : (loaded_start_idx + take) # type: ignore[misc] + ] # type: ignore[misc] + + # move indexing boundaries + boundary += shard_size + loaded_boundary += full_dim - extra + + return loader diff --git a/python/sglang/srt/layers/attention/merge_state.py b/python/sglang/srt/layers/attention/merge_state.py new file mode 100644 index 000000000..245418a9d --- /dev/null +++ b/python/sglang/srt/layers/attention/merge_state.py @@ -0,0 +1,46 @@ +from typing import Optional, Tuple + +import torch +from sgl_kernel import merge_state_v2 + +from sglang.srt.layers.attention.triton_ops.merge_state import merge_state_triton +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() + + +# Automatically fallback to the Triton kernel in some cases +# (e.g., for AMD GPUs, when the head dimension is not a multiple +# of 4 or 8, and in FP8 precision) +def _supported_dtypes(o: torch.Tensor) -> bool: + return o.dtype in [torch.float32, torch.half, torch.bfloat16] + + +def _supported_headdim(o: torch.Tensor) -> bool: + headdim = o.shape[2] # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] + if o.dtype == torch.float32: + return headdim % 4 == 0 + return headdim % 8 == 0 + + +def merge_state( + prefix_output: torch.Tensor, + prefix_lse: torch.Tensor, + suffix_output: torch.Tensor, + suffix_lse: torch.Tensor, + output: Optional[torch.Tensor] = None, + output_lse: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if ( + _is_cuda + and _supported_dtypes(prefix_output) + and _supported_headdim(prefix_output) + ): + return merge_state_v2( + prefix_output, prefix_lse, suffix_output, suffix_lse, output, output_lse + ) + else: + # Fallback to Triton kernel + return merge_state_triton( + prefix_output, prefix_lse, suffix_output, suffix_lse, output, output_lse + ) diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py new file mode 100644 index 000000000..06cfbd4ef --- /dev/null +++ b/python/sglang/srt/layers/attention/tbo_backend.py @@ -0,0 +1,260 @@ +from typing import TYPE_CHECKING, Callable, List, Optional, Union + +import torch + +from sglang.srt import two_batch_overlap +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode + + +class TboAttnBackend(AttentionBackend): + def __init__(self, primary: AttentionBackend, children: List[AttentionBackend]): + super().__init__() + self.primary = primary + self.children = children + + @classmethod + def init_new(cls, creator: Callable[[], AttentionBackend]): + return cls( + primary=creator(), + children=[creator() for _ in range(2)], + ) + + def init_forward_metadata(self, forward_batch: "ForwardBatch"): + self.primary.init_forward_metadata(forward_batch=forward_batch) + if forward_batch.tbo_children is not None: + for child, forward_batch_child in zip( + self.children, forward_batch.tbo_children, strict=True + ): + if forward_batch_child.batch_size > 0: + child.init_forward_metadata(forward_batch=forward_batch_child) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.primary.init_cuda_graph_state(max_bs=max_bs, max_num_tokens=max_num_tokens) + for item in self.children: + # TODO for children, maybe can provide *smaller* max_bs to optimize + item.init_cuda_graph_state(max_bs=max_bs, max_num_tokens=max_num_tokens) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: "ForwardMode", + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + self.primary.init_forward_metadata_capture_cuda_graph( + bs=bs, + num_tokens=num_tokens, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + encoder_lens=encoder_lens, + forward_mode=forward_mode, + spec_info=spec_info, + ) + + self._init_forward_metadata_cuda_graph_children( + fn_name="init_forward_metadata_capture_cuda_graph", + bs=bs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + encoder_lens=encoder_lens, + forward_mode=forward_mode, + spec_info=spec_info, + capture_num_tokens=num_tokens, + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: "ForwardMode", + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + ): + self.primary.init_forward_metadata_replay_cuda_graph( + bs=bs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_sum=seq_lens_sum, + encoder_lens=encoder_lens, + forward_mode=forward_mode, + spec_info=spec_info, + seq_lens_cpu=seq_lens_cpu, + ) + + self._init_forward_metadata_cuda_graph_children( + fn_name="init_forward_metadata_replay_cuda_graph", + bs=bs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + encoder_lens=encoder_lens, + forward_mode=forward_mode, + spec_info=spec_info, + replay_seq_lens_sum=seq_lens_sum, + replay_seq_lens_cpu=seq_lens_cpu, + ) + + def _init_forward_metadata_cuda_graph_children( + self, + fn_name: str, + # common args + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: "ForwardMode", + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + # capture args + capture_num_tokens: int = None, + # replay args + replay_seq_lens_sum: int = None, + replay_seq_lens_cpu: Optional[torch.Tensor] = None, + ): + token_num_per_seq = two_batch_overlap.get_token_num_per_seq( + forward_mode=forward_mode, spec_info=spec_info + ) + if fn_name == "init_forward_metadata_capture_cuda_graph": + assert ( + capture_num_tokens == bs * token_num_per_seq + ), "For target-verify or decode mode, num_tokens should be equal to token_num_per_seq * bs" + num_tokens = bs * token_num_per_seq + + tbo_split_seq_index, tbo_split_token_index = ( + two_batch_overlap.compute_split_indices_for_cuda_graph_replay( + forward_mode=forward_mode, + cuda_graph_num_tokens=num_tokens, + spec_info=spec_info, + ) + ) + + num_tokens_child_left = tbo_split_token_index + num_tokens_child_right = num_tokens - tbo_split_token_index + bs_child_left = tbo_split_seq_index + bs_child_right = bs - bs_child_left + + assert ( + num_tokens_child_left > 0 and num_tokens_child_right > 0 + ), f"{num_tokens_child_left=} {num_tokens_child_right=} {forward_mode=} {num_tokens=}" + + common_pre_split_args = dict( + fn_name=fn_name, + bs=bs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + encoder_lens=encoder_lens, + forward_mode=forward_mode, + spec_info=spec_info, + capture_num_tokens=capture_num_tokens, + replay_seq_lens_sum=replay_seq_lens_sum, + replay_seq_lens_cpu=replay_seq_lens_cpu, + ) + + args_left = _init_forward_metadata_cuda_graph_split( + output_bs=bs_child_left, + seq_slice=slice(None, tbo_split_seq_index), + **common_pre_split_args, + ) + args_right = _init_forward_metadata_cuda_graph_split( + output_bs=bs_child_right, + seq_slice=slice(tbo_split_seq_index, None), + **common_pre_split_args, + ) + + child_left, child_right = self.children + getattr(child_left, fn_name)(**args_left) + getattr(child_right, fn_name)(**args_right) + + def get_cuda_graph_seq_len_fill_value(self): + ans = self.primary.get_cuda_graph_seq_len_fill_value() + for child in self.children: + assert ans == child.get_cuda_graph_seq_len_fill_value() + return ans + + def forward_extend(self, *args, **kwargs): + return self.primary.forward_extend(*args, **kwargs) + + def forward_decode(self, *args, **kwargs): + return self.primary.forward_decode(*args, **kwargs) + + +def _init_forward_metadata_cuda_graph_split( + fn_name: str, + seq_slice: slice, + output_bs: int, + # common args + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: "ForwardMode", + spec_info: Optional[EagleVerifyInput], + # capture args + capture_num_tokens: int = None, + # replay args + replay_seq_lens_sum: int = None, + replay_seq_lens_cpu: Optional[torch.Tensor] = None, +): + token_num_per_seq = two_batch_overlap.get_token_num_per_seq( + forward_mode=forward_mode, spec_info=spec_info + ) + assert encoder_lens is None, "encoder_lens is not supported yet" + if spec_info is not None: + output_spec_info = two_batch_overlap.split_spec_info( + spec_info=spec_info, + start_seq_index=seq_slice.start if seq_slice.start is not None else 0, + end_seq_index=seq_slice.stop if seq_slice.stop is not None else bs, + start_token_index=( + seq_slice.start * token_num_per_seq + if seq_slice.start is not None + else 0 + ), + end_token_index=( + seq_slice.stop * token_num_per_seq + if seq_slice.stop is not None + else bs * token_num_per_seq + ), + ) + + else: + output_spec_info = None + ans = dict( + bs=output_bs, + req_pool_indices=req_pool_indices[seq_slice], + seq_lens=seq_lens[seq_slice], + # directly forward + forward_mode=forward_mode, + # ignore + encoder_lens=None, + spec_info=output_spec_info, + ) + + if fn_name == "init_forward_metadata_capture_cuda_graph": + assert ( + capture_num_tokens == bs * token_num_per_seq + ), "Only support num_tokens==bs * token_num_per_seq for target-verify or decode mode" + ans.update( + dict( + num_tokens=output_bs * token_num_per_seq, + ) + ) + elif fn_name == "init_forward_metadata_replay_cuda_graph": + output_seq_lens_cpu = replay_seq_lens_cpu[seq_slice] + ans.update( + dict( + seq_lens_sum=output_seq_lens_cpu.sum().item(), + seq_lens_cpu=output_seq_lens_cpu, + ) + ) + else: + raise NotImplementedError + + return ans diff --git a/python/sglang/srt/layers/attention/torch_native_backend.py b/python/sglang/srt/layers/attention/torch_native_backend.py new file mode 100644 index 000000000..6a67ea947 --- /dev/null +++ b/python/sglang/srt/layers/attention/torch_native_backend.py @@ -0,0 +1,276 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from torch.nn.functional import scaled_dot_product_attention + +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.radix_attention import AttentionType +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + + +class TorchNativeAttnBackend(AttentionBackend): + def __init__(self, model_runner: ModelRunner): + super().__init__() + self.forward_metadata = None + self.device = model_runner.device + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init the metadata for a forward pass.""" + pass + + def _run_sdpa_forward_extend( + self, + query: torch.Tensor, + output: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + extend_prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + scaling=None, + enable_gqa=False, + causal=False, + ): + """Run the extend forward by using torch native sdpa op. + + Args: + query: [num_tokens, num_heads, head_size] + output: [num_tokens, num_heads, head_size] + k_cache: [max_total_num_tokens, num_heads, head_size] + v_cache: [max_total_num_tokens, num_heads, head_size] + req_to_token: [max_num_reqs, max_context_len] + req_pool_indices: [num_seqs] + seq_lens: [num_seqs] + extend_prefix_lens: [num_seqs] + extend_seq_lens: [num_seqs] + scaling: float or None + enable_gqa: bool + causal: bool + + Returns: + output: [num_tokens, num_heads, head_size] + """ + + assert seq_lens.shape[0] == extend_prefix_lens.shape[0] + assert seq_lens.shape[0] == extend_seq_lens.shape[0] + + # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size] + query = query.movedim(0, query.dim() - 2) + + start_q, start_kv = 0, 0 + for seq_idx in range(seq_lens.shape[0]): + # TODO: this loop process a sequence per iter, this is inefficient. + # Need optimize the performance later. + + extend_seq_len_q = extend_seq_lens[seq_idx] + prefill_seq_len_q = extend_prefix_lens[seq_idx] + + seq_len_kv = seq_lens[seq_idx] + end_q = start_q + extend_seq_len_q + end_kv = start_kv + seq_len_kv + + per_req_query = query[:, start_q:end_q, :] + per_req_query_redudant = torch.empty( + (per_req_query.shape[0], seq_len_kv, per_req_query.shape[2]), + dtype=per_req_query.dtype, + device=per_req_query.device, + ) + + per_req_query_redudant[:, prefill_seq_len_q:, :] = per_req_query + + # get key and value from cache. per_req_tokens contains the kv cache + # index for each token in the sequence. + req_pool_idx = req_pool_indices[seq_idx] + per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv] + per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2) + per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2) + + per_req_out_redudant = ( + scaled_dot_product_attention( + per_req_query_redudant.unsqueeze(0), + per_req_key.unsqueeze(0), + per_req_value.unsqueeze(0), + enable_gqa=enable_gqa, + scale=scaling, + is_causal=causal, + ) + .squeeze(0) + .movedim(query.dim() - 2, 0) + ) + output[start_q:end_q, :, :] = per_req_out_redudant[prefill_seq_len_q:, :, :] + start_q, start_kv = end_q, end_kv + return output + + def _run_sdpa_forward_decode( + self, + query: torch.Tensor, + output: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + scaling=None, + enable_gqa=False, + causal=False, + ): + """Run the decode forward by using torch native sdpa op. + + Args: + query: [num_tokens, num_heads, head_size] + output: [num_tokens, num_heads, head_size] + k_cache: [max_total_num_tokens, num_heads, head_size] + v_cache: [max_total_num_tokens, num_heads, head_size] + req_to_token: [max_num_reqs, max_context_len] + req_pool_indices: [num_seqs] + seq_lens: [num_seqs] + scaling: float or None + enable_gqa: bool + causal: bool + + Returns: + output: [num_tokens, num_heads, head_size] + """ + + # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size] + query = query.movedim(0, query.dim() - 2) + + start_q, start_kv = 0, 0 + for seq_idx in range(seq_lens.shape[0]): + # TODO: this loop process a sequence per iter, this is inefficient. + # Need optimize the performance later. + + seq_len_q = 1 + seq_len_kv = seq_lens[seq_idx] + end_q = start_q + seq_len_q + end_kv = start_kv + seq_len_kv + + per_req_query = query[:, start_q:end_q, :] + + # get key and value from cache. per_req_tokens contains the kv cache + # index for each token in the sequence. + req_pool_idx = req_pool_indices[seq_idx] + per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv] + per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2) + per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2) + + per_req_out = ( + scaled_dot_product_attention( + per_req_query.unsqueeze(0), + per_req_key.unsqueeze(0), + per_req_value.unsqueeze(0), + enable_gqa=enable_gqa, + scale=scaling, + is_causal=causal, + ) + .squeeze(0) + .movedim(query.dim() - 2, 0) + ) + output[start_q:end_q, :, :] = per_req_out + start_q, start_kv = end_q, end_kv + + return output + + def forward_extend( + self, + q, + k, + v, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if layer.is_cross_attention: + cache_loc = forward_batch.encoder_out_cache_loc + else: + cache_loc = forward_batch.out_cache_loc + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) + + use_gqa = layer.tp_q_head_num != layer.tp_k_head_num + + q_ = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) + o_ = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) + + causal = True + if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY: + causal = False + + self._run_sdpa_forward_extend( + q_, + o_, + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + forward_batch.req_to_token_pool.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.extend_prefix_lens, + forward_batch.extend_seq_lens, + scaling=layer.scaling, + enable_gqa=use_gqa, + causal=causal, + ) + return o + + def forward_decode( + self, + q, + k, + v, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + # During torch.compile, there is a bug in rotary_emb that causes the + # output value to have a 3D tensor shape. This reshapes the output correctly. + q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim) + + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if layer.is_cross_attention: + cache_loc = forward_batch.encoder_out_cache_loc + else: + cache_loc = forward_batch.out_cache_loc + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) + + use_gqa = layer.tp_q_head_num != layer.tp_k_head_num + + q_ = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) + o_ = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) + + self._run_sdpa_forward_decode( + q_, + o_, + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + forward_batch.req_to_token_pool.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + scaling=layer.scaling, + enable_gqa=use_gqa, + causal=False, + ) + + return o + + def support_triton(self): + return False diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py new file mode 100644 index 000000000..26241d849 --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -0,0 +1,1090 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Union + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.layers.radix_attention import AttentionType +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils import get_bool_env_var, get_device_core_count, next_power_of_2 + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput + + +def logit_capping_mod(logit_capping_method, logit_cap): + # positive logit_cap -> tanh cap + if logit_capping_method == "tanh": + return logit_cap + else: + raise ValueError() + + +@dataclass +class ForwardMetadata: + attn_logits: torch.Tensor + attn_lse: torch.Tensor + max_extend_len: int + num_kv_splits: torch.Tensor + kv_indptr: torch.Tensor + kv_indices: torch.Tensor + qo_indptr: torch.Tensor + custom_mask: torch.Tensor + mask_indptr: torch.Tensor + # Sliding window + window_kv_indptr: torch.Tensor + window_kv_indices: torch.Tensor + window_num_kv_splits: torch.Tensor + window_kv_offsets: torch.Tensor + + +class TritonAttnBackend(AttentionBackend): + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + kv_indptr_buf: Optional[torch.Tensor] = None, + ): + # Lazy import to avoid the initialization of cuda context + from sglang.srt.layers.attention.triton_ops.decode_attention import ( + decode_attention_fwd, + ) + from sglang.srt.layers.attention.triton_ops.extend_attention import ( + extend_attention_fwd, + ) + + super().__init__() + + self.decode_attention_fwd = torch.compiler.disable(decode_attention_fwd) + self.extend_attention_fwd = torch.compiler.disable(extend_attention_fwd) + + # Parse args + self.skip_prefill = skip_prefill + max_bs = model_runner.req_to_token_pool.size + self.sliding_window_size = model_runner.sliding_window_size + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.token_to_kv_pool_allocator = model_runner.token_to_kv_pool_allocator + self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens + self.speculative_num_steps = model_runner.server_args.speculative_num_steps + self.num_head = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.num_kv_head = model_runner.model_config.get_num_kv_heads( + get_attention_tp_size() + ) + self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1] + self.max_context_len = model_runner.model_config.context_len + self.device = model_runner.device + self.device_core_count = get_device_core_count(model_runner.gpu_id) + self.static_kv_splits = get_bool_env_var( + "SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS", "false" + ) + self.max_kv_splits = model_runner.server_args.triton_attention_num_kv_splits + + # Check arguments + assert not ( + model_runner.sliding_window_size is not None + and model_runner.model_config.is_encoder_decoder + ), "Sliding window and cross attention are not supported together" + + # Initialize buffers + # TODO(Jianan Ji): Make sure it behaves as expected when kv_indptr_buf is provided and sliding window is enabled + if kv_indptr_buf is None: + self.kv_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + else: + self.kv_indptr = kv_indptr_buf + + # If sliding window is enabled, we might need two sets of buffers + # because of interleaved attention types (e.g. for Gemma3) + self.window_kv_indptr = None + if self.sliding_window_size is not None and self.sliding_window_size > 0: + if kv_indptr_buf is None: + self.window_kv_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + else: + # When provided a buffer, create a clone for the second buffer + self.window_kv_indptr = torch.zeros_like(kv_indptr_buf) + + if not self.skip_prefill: + self.qo_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + + self.mask_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int64, device=model_runner.device + ) + + # Initialize forward metadata + self.forward_metadata: ForwardMetadata = None + + def get_num_kv_splits( + self, + num_kv_splits: torch.Tensor, + seq_lens: torch.Tensor, + ): + num_token, num_seq = num_kv_splits.shape[0], seq_lens.shape[0] + # NOTE(alcanderian): Considering speculative_decodeing, + # num_kv_splits.shape[0] will be topk * real_num_token. + # And the real_num_token is num_seq in decoding phase. + num_group = num_token // num_seq + + assert ( + num_group * num_seq == num_token + ), f"num_seq({num_seq}), num_token({num_token}), something goes wrong!" + + if self.static_kv_splits or self.device_core_count <= 0: + num_kv_splits.fill_(self.max_kv_splits) + return + + if num_seq < 256: + SCHEDULE_SEQ = 256 + else: + SCHEDULE_SEQ = triton.next_power_of_2(num_seq) + + get_num_kv_splits_triton[(1,)]( + num_kv_splits, + seq_lens, + num_seq, + num_group, + self.num_head, + self.num_kv_head, + self.max_kv_splits, + self.device_core_count, + MAX_NUM_SEQ=SCHEDULE_SEQ, + ) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init auxiliary variables for triton attention backend.""" + + bs = forward_batch.batch_size + kv_indptr = self.kv_indptr + window_kv_indptr = self.window_kv_indptr + window_kv_indices = None + window_num_kv_splits = None + window_kv_offsets = None + spec_info = forward_batch.spec_info + + if forward_batch.forward_mode.is_decode_or_idle(): + if spec_info is None: + kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + forward_batch.seq_lens_sum, dtype=torch.int64, device=self.device + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + # Sliding window + if ( + self.sliding_window_size is not None + and self.sliding_window_size > 0 + ): + window_kv_indptr, window_kv_indices, window_kv_lens, _ = ( + update_sliding_window_buffer( + self.window_kv_indptr, + self.req_to_token, + self.sliding_window_size, + forward_batch.seq_lens, + forward_batch.req_pool_indices, + bs, + self.device, + self.token_to_kv_pool_allocator, + ) + ) + window_num_kv_splits = torch.empty( + (bs,), dtype=torch.int32, device=self.device + ) + self.get_num_kv_splits(window_num_kv_splits, window_kv_lens) + else: + kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices + bs = kv_indptr.shape[0] - 1 + + attn_logits = torch.empty( + (bs, self.num_head, self.max_kv_splits, self.v_head_dim), + dtype=torch.float32, + device=self.device, + ) + attn_lse = torch.empty( + (bs, self.num_head, self.max_kv_splits), + dtype=torch.float32, + device=self.device, + ) + num_kv_splits = torch.empty((bs,), dtype=torch.int32, device=self.device) + self.get_num_kv_splits(num_kv_splits, forward_batch.seq_lens) + + qo_indptr = None + custom_mask = None + mask_indptr = None + max_extend_len = None + elif forward_batch.forward_mode.is_target_verify(): + bs = len(forward_batch.req_pool_indices) + qo_indptr = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + # Different with flashinfer kv_indptr and kv_indices construction + kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + kv_indptr[-1], dtype=torch.int64, device=self.device + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + + if self.sliding_window_size is not None and self.sliding_window_size > 0: + # window_kv_offsets is used to calculate the start position in custom mask + ( + window_kv_indptr, + window_kv_indices, + window_kv_lens, + window_kv_offsets, + ) = update_sliding_window_buffer( + self.window_kv_indptr, + self.req_to_token, + self.sliding_window_size, + forward_batch.seq_lens, + forward_batch.req_pool_indices, + bs, + self.device, + self.token_to_kv_pool_allocator, + ) + + custom_mask = spec_info.custom_mask + seq_mask_len = self.num_draft_tokens * ( + forward_batch.seq_lens + self.num_draft_tokens + ) + mask_indptr = self.mask_indptr + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0) + mask_indptr = mask_indptr[: bs + 1] + max_extend_len = self.num_draft_tokens + num_kv_splits = None + attn_logits = None + attn_lse = None + + elif forward_batch.forward_mode.is_draft_extend(): + kv_indices, kv_indptr, qo_indptr, custom_mask = ( + spec_info.generate_attn_arg_prefill( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + None, + self.req_to_token, + ) + ) + kv_indices = kv_indices.to(torch.int64) + mask_indptr = None + # TODO(FIXME): This will trigger an invalid Eagle tree when using + # `max(spec_info.accept_length_cpu)`. + # It might have been forgotten to update somewhere. + max_extend_len = torch.max(spec_info.accept_length).item() + num_kv_splits = None + attn_logits = None + attn_lse = None + else: + kv_indptr[1 : bs + 1] = torch.cumsum( + forward_batch.extend_prefix_lens, dim=0 + ) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + forward_batch.extend_prefix_lens.sum().item(), + dtype=torch.int64, + device=self.device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.extend_prefix_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + # Sliding window + if self.sliding_window_size is not None and self.sliding_window_size > 0: + window_kv_indptr, window_kv_indices, _, _ = ( + update_sliding_window_buffer( + self.window_kv_indptr, + self.req_to_token, + self.sliding_window_size, + forward_batch.extend_prefix_lens, + forward_batch.req_pool_indices, + bs, + self.device, + self.token_to_kv_pool_allocator, + ) + ) + + qo_indptr = self.qo_indptr + qo_indptr[1 : bs + 1] = torch.cumsum(forward_batch.extend_seq_lens, dim=0) + qo_indptr = qo_indptr[: bs + 1] + custom_mask = None + mask_indptr = None + attn_logits = None + attn_lse = None + max_extend_len = max(forward_batch.extend_seq_lens_cpu) + num_kv_splits = None + + self.forward_metadata = ForwardMetadata( + attn_logits, + attn_lse, + max_extend_len, + num_kv_splits, + kv_indptr, + kv_indices, + qo_indptr, + custom_mask, + mask_indptr, + window_kv_indptr, + window_kv_indices, + window_num_kv_splits, + window_kv_offsets, + ) + + def init_cuda_graph_state( + self, + max_bs: int, + max_num_tokens: int, + kv_indices_buf: Optional[torch.Tensor] = None, + ): + self.cuda_graph_attn_logits = torch.zeros( + (max_num_tokens, self.num_head, self.max_kv_splits, self.v_head_dim), + dtype=torch.float32, + device=self.device, + ) + self.cuda_graph_attn_lse = torch.zeros( + (max_num_tokens, self.num_head, self.max_kv_splits), + dtype=torch.float32, + device=self.device, + ) + self.cuda_graph_num_kv_splits = torch.full( + (max_num_tokens,), self.max_kv_splits, dtype=torch.int32, device=self.device + ) + if kv_indices_buf is None: + self.cuda_graph_kv_indices = torch.zeros( + (max_num_tokens * self.max_context_len), + dtype=torch.int64, + device=self.device, + ) + else: + self.cuda_graph_kv_indices = kv_indices_buf + + if not self.skip_prefill: + self.cuda_graph_custom_mask = torch.zeros( + (max_num_tokens * self.max_context_len), + dtype=torch.uint8, + device=self.device, + ) + + if self.sliding_window_size is not None and self.sliding_window_size > 0: + if kv_indices_buf is None: + self.cuda_graph_window_kv_indices = torch.zeros( + (max_num_tokens * self.sliding_window_size), + dtype=torch.int64, + device=self.device, + ) + else: + self.cuda_graph_window_kv_indices = torch.zeros_like(kv_indices_buf) + + self.cuda_graph_window_num_kv_splits = torch.full( + (max_num_tokens,), + self.max_kv_splits, + dtype=torch.int32, + device=self.device, + ) + + self.cuda_graph_window_kv_offsets = torch.zeros( + (max_bs,), + dtype=torch.int32, + device=self.device, + ) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + assert encoder_lens is None, "Not supported" + window_kv_indptr = self.window_kv_indptr + window_kv_indices = None + window_num_kv_splits = None + window_kv_offsets = None + + if forward_mode.is_decode_or_idle(): + if spec_info is None: + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + if ( + self.sliding_window_size is not None + and self.sliding_window_size > 0 + ): + window_kv_indices = self.cuda_graph_window_kv_indices + window_num_kv_splits = self.cuda_graph_window_num_kv_splits + window_kv_indptr, window_kv_indices, _, _ = ( + update_sliding_window_buffer_cuda_graph( + self.window_kv_indptr, + window_kv_indices, + self.req_to_token, + self.sliding_window_size, + seq_lens[:bs], + req_pool_indices, + bs, + self.token_to_kv_pool_allocator, + ) + ) + else: + kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices + + attn_logits = self.cuda_graph_attn_logits + attn_lse = self.cuda_graph_attn_lse + max_extend_len = None + num_kv_splits = self.cuda_graph_num_kv_splits + qo_indptr = None + custom_mask = None + mask_indptr = None + elif forward_mode.is_target_verify(): + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + + if self.sliding_window_size is not None and self.sliding_window_size > 0: + window_kv_indices = self.cuda_graph_window_kv_indices + window_num_kv_splits = self.cuda_graph_window_num_kv_splits + window_kv_offsets = self.cuda_graph_window_kv_offsets + window_kv_indptr, window_kv_indices, _, window_kv_offsets[:bs] = ( + update_sliding_window_buffer_cuda_graph( + self.window_kv_indptr, + window_kv_indices, + self.req_to_token, + self.sliding_window_size, + seq_lens[:bs], + req_pool_indices, + bs, + self.token_to_kv_pool_allocator, + ) + ) + + custom_mask = self.cuda_graph_custom_mask + custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask + seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens) + mask_indptr = self.mask_indptr[: bs + 1] + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0) + max_extend_len = self.num_draft_tokens + num_kv_splits = None + attn_logits = None + attn_lse = None + elif forward_mode.is_draft_extend(): + num_tokens_per_bs = self.speculative_num_steps + 1 + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + bs * num_tokens_per_bs + 1, + step=num_tokens_per_bs, + dtype=torch.int32, + device=self.device, + ) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + custom_mask = None + mask_indptr = None + max_extend_len = num_tokens_per_bs + num_kv_splits = None + attn_logits = None + attn_lse = None + else: + raise ValueError( + f"Invalid forward mode: {forward_mode=} for CUDA Graph capture." + ) + + self.forward_metadata = ForwardMetadata( + attn_logits, + attn_lse, + max_extend_len, + num_kv_splits, + kv_indptr, + kv_indices, + qo_indptr, + custom_mask, + mask_indptr, + window_kv_indptr, + window_kv_indices, + window_num_kv_splits, + window_kv_offsets, + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + ): + # NOTE: encoder_lens expected to be zeros or None + if forward_mode.is_decode_or_idle(): + # Update kv_indptr, kv_indices + kv_indptr = self.kv_indptr + kv_indices = self.cuda_graph_kv_indices + num_kv_splits = self.cuda_graph_num_kv_splits + if spec_info is None: + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens[:bs], dim=0) + kv_indptr = kv_indptr[: bs + 1] + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices[:bs], + seq_lens[:bs], + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + num_token = bs + if ( + self.sliding_window_size is not None + and self.sliding_window_size > 0 + ): + window_num_kv_splits = self.cuda_graph_window_num_kv_splits + window_kv_indices = self.cuda_graph_window_kv_indices + _, _, window_kv_lens, _ = update_sliding_window_buffer_cuda_graph( + self.window_kv_indptr, + window_kv_indices, + self.req_to_token, + self.sliding_window_size, + seq_lens[:bs], + req_pool_indices[:bs], + bs, + self.token_to_kv_pool_allocator, + ) + self.get_num_kv_splits( + window_num_kv_splits[:num_token], window_kv_lens[:bs] + ) + + else: + kv_indptr[: spec_info.kv_indptr.shape[0]] = spec_info.kv_indptr + kv_indices[: spec_info.kv_indices.shape[0]] = spec_info.kv_indices + num_token = spec_info.kv_indptr.shape[0] - 1 + self.get_num_kv_splits(num_kv_splits[:num_token], seq_lens[:bs]) + + elif forward_mode.is_target_verify(): + # Update qo_indptr, kv_indptr, kv_indices, custom_mask, mask_indptr + bs = len(req_pool_indices) + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + if self.sliding_window_size is not None and self.sliding_window_size > 0: + window_num_kv_splits = self.cuda_graph_window_num_kv_splits + window_kv_indices = self.cuda_graph_window_kv_indices + window_kv_offsets = self.cuda_graph_window_kv_offsets + _, _, window_kv_lens, window_kv_offsets[:bs] = ( + update_sliding_window_buffer_cuda_graph( + self.window_kv_indptr, + window_kv_indices, + self.req_to_token, + self.sliding_window_size, + seq_lens[:bs], + req_pool_indices, + bs, + self.token_to_kv_pool_allocator, + ) + ) + custom_mask = self.cuda_graph_custom_mask + custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask + seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens) + mask_indptr = self.mask_indptr[: bs + 1] + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0) + elif forward_mode.is_draft_extend(): + seq_lens = seq_lens[:bs] + accept_lens = spec_info.accept_length[:bs] + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[1 : bs + 1] = torch.cumsum(accept_lens, dim=0) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + else: + raise ValueError( + f"Invalid forward mode: {forward_mode=} for CUDA Graph replay." + ) + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + sinks=None, + ): + # TODO: reuse the buffer across layers + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + + logits_soft_cap = logit_capping_mod(layer.logit_capping_method, layer.logit_cap) + + causal = True + if layer.attn_type == AttentionType.ENCODER_ONLY: + causal = False + + if layer.sliding_window_size is not None and layer.sliding_window_size > -1: + sliding_window_size = ( + layer.sliding_window_size + ) # Needed for sliding window mask + kv_indptr = self.forward_metadata.window_kv_indptr + kv_indices = self.forward_metadata.window_kv_indices + window_kv_offsets = self.forward_metadata.window_kv_offsets + else: + sliding_window_size = -1 + kv_indptr = self.forward_metadata.kv_indptr + kv_indices = self.forward_metadata.kv_indices + window_kv_offsets = None + + self.extend_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + k.contiguous(), + v.contiguous(), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + self.forward_metadata.qo_indptr, + kv_indptr, + kv_indices, + self.forward_metadata.custom_mask, + causal, + self.forward_metadata.mask_indptr, + self.forward_metadata.max_extend_len, + layer.scaling, + logit_cap=logits_soft_cap, + sliding_window_size=sliding_window_size, + sinks=sinks, + window_kv_offsets=window_kv_offsets, + xai_temperature_len=layer.xai_temperature_len, + ) + return o + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + sinks=None, + ): + # During torch.compile, there is a bug in rotary_emb that causes the + # output value to have a 3D tensor shape. This reshapes the output correctly. + q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim) + + # TODO: reuse the buffer across layers + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + logits_soft_cap = logit_capping_mod(layer.logit_capping_method, layer.logit_cap) + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + + if layer.sliding_window_size is not None and layer.sliding_window_size > -1: + kv_indptr = self.forward_metadata.window_kv_indptr + kv_indices = self.forward_metadata.window_kv_indices + else: + kv_indptr = self.forward_metadata.kv_indptr + kv_indices = self.forward_metadata.kv_indices + + self.decode_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + kv_indptr, + kv_indices, + self.forward_metadata.attn_logits, + self.forward_metadata.attn_lse, + self.forward_metadata.num_kv_splits, + self.max_kv_splits, + layer.scaling, + logit_cap=logits_soft_cap, + sinks=sinks, + xai_temperature_len=layer.xai_temperature_len, + ) + return o + + +class TritonMultiStepDraftBackend: + """ + Wrap multiple triton attention backends as one for multiple consecutive + draft decoding steps. + """ + + def __init__( + self, + model_runner: ModelRunner, + topk: int, + speculative_num_steps: int, + ): + from sglang.srt.speculative.eagle_utils import generate_draft_decode_kv_indices + + self.topk = topk + self.speculative_num_steps = speculative_num_steps + self.generate_draft_decode_kv_indices = generate_draft_decode_kv_indices + max_bs = model_runner.req_to_token_pool.size * self.topk + self.kv_indptr = torch.zeros( + ( + self.speculative_num_steps, + max_bs + 1, + ), + dtype=torch.int32, + device=model_runner.device, + ) + self.attn_backends = [] + for i in range(self.speculative_num_steps): + self.attn_backends.append( + TritonAttnBackend( + model_runner, + skip_prefill=True, + kv_indptr_buf=self.kv_indptr[i], + ) + ) + self.max_context_len = self.attn_backends[0].max_context_len + self.num_head = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.device = model_runner.device + # Cached variables for generate_draft_decode_kv_indices + self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1] + self.page_size = model_runner.server_args.page_size + + def common_template( + self, forward_batch: ForwardBatch, kv_indices_buffer: torch.Tensor, call_fn: int + ): + num_seqs = forward_batch.batch_size + bs = self.topk * num_seqs + seq_lens_sum = forward_batch.seq_lens_sum + + self.generate_draft_decode_kv_indices[ + (self.speculative_num_steps, num_seqs, self.topk) + ]( + forward_batch.req_pool_indices, + forward_batch.req_to_token_pool.req_to_token, + forward_batch.seq_lens, + kv_indices_buffer, + self.kv_indptr, + forward_batch.positions, + self.pool_len, + kv_indices_buffer.shape[1], + self.kv_indptr.shape[1], + next_power_of_2(num_seqs), + next_power_of_2(self.speculative_num_steps), + next_power_of_2(bs), + self.page_size, + ) + + for i in range(self.speculative_num_steps): + forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1] + forward_batch.spec_info.kv_indices = kv_indices_buffer[i][ + : seq_lens_sum * self.topk + bs * (i + 1) + ] + call_fn(i, forward_batch) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + kv_indices = torch.empty( + ( + self.speculative_num_steps, + forward_batch.batch_size * self.topk * self.max_context_len, + ), + dtype=torch.int64, + device=self.device, + ) + + def call_fn(i, forward_batch): + forward_batch.spec_info.kv_indptr = ( + forward_batch.spec_info.kv_indptr.clone() + ) + forward_batch.spec_info.kv_indices = ( + forward_batch.spec_info.kv_indices.clone() + ) + self.attn_backends[i].init_forward_metadata(forward_batch) + + self.common_template(forward_batch, kv_indices, call_fn) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.cuda_graph_kv_indices = torch.zeros( + (self.speculative_num_steps, max_num_tokens * self.max_context_len), + dtype=torch.int64, + device=self.device, + ) + for i in range(self.speculative_num_steps): + self.attn_backends[i].init_cuda_graph_state( + max_bs, max_num_tokens, kv_indices_buf=self.cuda_graph_kv_indices[i] + ) + + def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_capture_cuda_graph( + forward_batch.batch_size, + forward_batch.batch_size * self.topk, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + ) + + self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn) + + def init_forward_metadata_replay_cuda_graph( + self, forward_batch: ForwardBatch, bs: int + ): + def call_fn(i, forward_batch): + self.attn_backends[i].init_forward_metadata_replay_cuda_graph( + bs, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + seq_lens_sum=-1, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + seq_lens_cpu=None, + ) + + self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn) + + +@triton.jit +def get_num_kv_splits_triton( + num_kv_splits_ptr, + seq_lens_ptr, + num_seq, + num_group, + num_head, + num_kv_head, + max_kv_splits, + device_core_count, + MAX_NUM_SEQ: tl.constexpr, +): + # TODO: this method is tunable, we need more online serving data to tune it + offs_seq = tl.arange(0, MAX_NUM_SEQ) + mask_seq = offs_seq < num_seq + + seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=0) + max_seq_len = tl.max(seq_lens) + seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=max_seq_len) + min_seq_len = tl.min(seq_lens) + if max_seq_len * 8 < min_seq_len * 10: + min_seq_len = max_seq_len + max_kv_splits_1 = tl.minimum(tl.cdiv(max_seq_len, min_seq_len), max_kv_splits) + kv_chunk_size_1 = tl.cdiv(max_seq_len, max_kv_splits_1) + + # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually + ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0 + ext_device_core_count = tl.cast( + device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32 + ) + block_h, num_kv_group = 16, num_head // num_kv_head + if num_kv_group == 1: + token_grid = num_seq * num_group * num_head + else: + # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd + block_h = tl.minimum(block_h, num_kv_group) + token_grid = num_seq * num_group * tl.cdiv(num_head, block_h) + max_kv_splits_2 = tl.minimum( + tl.cdiv(ext_device_core_count, token_grid), max_kv_splits + ) + kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2) + + num_kv_splits = tl.maximum( + tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2) + ) + + offs_token = offs_seq * num_group + mask_token = offs_token < num_seq * num_group + for i in range(0, num_group): + tl.store(num_kv_splits_ptr + i + offs_token, num_kv_splits, mask=mask_token) + + +def update_sliding_window_buffer( + window_kv_indptr, + req_to_token, + sliding_window_size, + seq_lens, + req_pool_indices, + bs, + device, + token_to_kv_pool_allocator=None, +): + window_kv_lens = torch.minimum( + seq_lens, + torch.tensor(sliding_window_size), + ) + window_kv_indptr[1 : bs + 1] = torch.cumsum(window_kv_lens, dim=0) + window_kv_indptr = window_kv_indptr[: bs + 1] + window_kv_indices = torch.empty( + window_kv_indptr[-1], dtype=torch.int64, device=device + ) + window_kv_start_idx = seq_lens - window_kv_lens + create_flashinfer_kv_indices_triton[(bs,)]( + req_to_token, + req_pool_indices, + window_kv_lens, + window_kv_indptr, + window_kv_start_idx, + window_kv_indices, + req_to_token.stride(0), + ) + # full to swa index mapping + if hasattr(token_to_kv_pool_allocator, "translate_loc_from_full_to_swa"): + kv_last_index = window_kv_indptr[-1] + window_kv_indices[:kv_last_index] = ( + token_to_kv_pool_allocator.translate_loc_from_full_to_swa( + window_kv_indices[:kv_last_index] + ) + ) + return window_kv_indptr, window_kv_indices, window_kv_lens, window_kv_start_idx + + +def update_sliding_window_buffer_cuda_graph( + window_kv_indptr, + window_kv_indices, + req_to_token, + sliding_window_size, + seq_lens, + req_pool_indices, + bs, + token_to_kv_pool_allocator=None, +): + window_kv_lens = torch.minimum( + seq_lens, + torch.tensor(sliding_window_size), + ) + window_kv_indptr[1 : bs + 1] = torch.cumsum(window_kv_lens, dim=0) + window_kv_indptr = window_kv_indptr[: bs + 1] + window_kv_start_idx = seq_lens - window_kv_lens + create_flashinfer_kv_indices_triton[(bs,)]( + req_to_token, + req_pool_indices, + window_kv_lens, + window_kv_indptr, + window_kv_start_idx, + window_kv_indices, + req_to_token.stride(0), + ) + # full to swa index mapping + if hasattr(token_to_kv_pool_allocator, "translate_loc_from_full_to_swa"): + kv_last_index = window_kv_indptr[-1] + window_kv_indices[:kv_last_index] = ( + token_to_kv_pool_allocator.translate_loc_from_full_to_swa( + window_kv_indices[:kv_last_index] + ) + ) + return window_kv_indptr, window_kv_indices, window_kv_lens, window_kv_start_idx diff --git a/python/sglang/srt/layers/attention/triton_ops/decode_attention.py b/python/sglang/srt/layers/attention/triton_ops/decode_attention.py new file mode 100644 index 000000000..1ba5d463d --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/decode_attention.py @@ -0,0 +1,776 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +""" +Memory-efficient attention for decoding. +It supports page size = 1. +""" + +# Adapted from +# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py +# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py + +import logging + +import triton +import triton.language as tl + +from sglang.srt.utils import is_hip + +_is_hip = is_hip() + +logger = logging.getLogger(__name__) + + +_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, + xai_temperature_len: 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) + + if xai_temperature_len > 0: + offs_qidx = cur_batch_seq_len - 1 + xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len)) + _qtemp = tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale + xai_temperature_reg = tl.where(offs_qidx > xai_temperature_len, _qtemp, 1.0) + + 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) + + if xai_temperature_len > 0: + qk *= xai_temperature_reg + + 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, + xai_temperature_len=-1, +): + 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 = q.shape[0], 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, + xai_temperature_len=xai_temperature_len, + 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, + xai_temperature_len: 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) + + 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) + + if xai_temperature_len > 0: + offs_qidx = cur_batch_seq_len - 1 + xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len)) + _qtemp = tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale + xai_temperature_reg = tl.where(offs_qidx > xai_temperature_len, _qtemp, 1.0) + + 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) + + if xai_temperature_len > 0: + qk *= xai_temperature_reg[:, None] + + 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, + xai_temperature_len=-1, +): + 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 = q.shape[0], 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, + xai_temperature_len=xai_temperature_len, + 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, + xai_temperature_len=-1, +): + _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, + xai_temperature_len, + ) + _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, + xai_temperature_len=-1, +): + _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, + xai_temperature_len, + ) + _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, + xai_temperature_len=-1, +): + 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, + xai_temperature_len=xai_temperature_len, + ) + 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, + xai_temperature_len=xai_temperature_len, + ) diff --git a/python/sglang/srt/layers/attention/triton_ops/double_sparsity_attention.py b/python/sglang/srt/layers/attention/triton_ops/double_sparsity_attention.py new file mode 100644 index 000000000..72e0bfe78 --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/double_sparsity_attention.py @@ -0,0 +1,1106 @@ +import torch +import triton +import triton.language as tl + +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.utils import is_cuda, is_hip + +_is_cuda = is_cuda() +if _is_cuda: + CUDA_CAPABILITY = torch.cuda.get_device_capability() + +_is_hip = is_hip() + +if global_server_args_dict.get("attention_reduce_in_fp32", False): + REDUCE_TRITON_TYPE = tl.float32 + REDUCE_TORCH_TYPE = torch.float32 +else: + REDUCE_TRITON_TYPE = tl.float16 + REDUCE_TORCH_TYPE = torch.float16 + + +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def _fwd_kernel_flash_decode_stage1( + Q, + K, + V, + sm_scale, + Req_to_tokens, + B_req_idx, + B_Seqlen, + Mid_O, # [batch, head, seq_block_num, head_dim] + Mid_O_LogExpSum, # [batch, head, seq_block_num] + stride_req_to_tokens_b, + stride_req_to_tokens_s, + stride_qbs, + stride_qh, + stride_qd, + stride_kbs, + stride_kh, + stride_kd, + stride_vbs, + stride_vh, + stride_vd, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_mid_od, + stride_mid_o_eb, + stride_mid_o_eh, + stride_mid_o_es, + gqa_group_size, + BLOCK_SEQ: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + seq_start_block = tl.program_id(2) + cur_kv_head = cur_head // gqa_group_size + + offs_d = tl.arange(0, BLOCK_DMODEL) + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_req_idx = tl.load(B_req_idx + cur_batch) + cur_batch_start_index = seq_start_block * BLOCK_SEQ + cur_batch_end_index = tl.minimum( + cur_batch_seq_len, cur_batch_start_index + BLOCK_SEQ + ) + + off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d + + block_n_size = ( + tl.where( + cur_batch_end_index - cur_batch_start_index <= 0, + 0, + cur_batch_end_index - cur_batch_start_index + BLOCK_N - 1, + ) + // BLOCK_N + ) + + offs_n = cur_batch_start_index + tl.arange(0, BLOCK_N) + + q = tl.load(Q + off_q) + + sum_exp = 0.0 + max_logic = -float("inf") + acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32) + + for start_n in range(0, block_n_size, 1): + offs_n_new = start_n * BLOCK_N + offs_n + k_loc = tl.load( + Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_n_new, + mask=offs_n_new < cur_batch_end_index, + other=0, + ) + off_k = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :] + k = tl.load( + K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0 + ) + att_value = tl.sum(q[None, :] * k, 1) + att_value *= sm_scale + att_value = tl.where(offs_n_new < cur_batch_end_index, att_value, float("-inf")) + v = tl.load( + V + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0 + ) + + cur_max_logic = tl.max(att_value, axis=0) + new_max_logic = tl.maximum(cur_max_logic, max_logic) + + exp_logic = tl.exp(att_value - new_max_logic) + logic_scale = tl.exp(max_logic - new_max_logic) + acc *= logic_scale + acc += tl.sum(exp_logic[:, None] * v, axis=0) + + sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=0) + max_logic = new_max_logic + + need_store = tl.where(block_n_size == 0, 0, 1) + for _ in range(0, need_store, 1): + off_mid_o = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + seq_start_block * stride_mid_os + + offs_d + ) + off_mid_o_logexpsum = ( + cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh + seq_start_block + ) + tl.store(Mid_O + off_mid_o, acc / sum_exp) + tl.store(Mid_O_LogExpSum + off_mid_o_logexpsum, max_logic + tl.log(sum_exp)) + return + + +@triton.jit +def _fwd_kernel_flash_decode_stage2( + B_Seqlen, + Mid_O, # [batch, head, seq_block_num, head_dim] + Mid_O_LogExpSum, # [batch, head, seq_block_num] + O, # [batch, head, head_dim] + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_mid_od, + stride_mid_o_eb, + stride_mid_o_eh, + stride_mid_o_es, + stride_obs, + stride_oh, + stride_od, + BLOCK_SEQ: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + + offs_d = tl.arange(0, BLOCK_DMODEL) + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + + block_n_size = ( + tl.where(cur_batch_seq_len <= 0, 0, cur_batch_seq_len + BLOCK_SEQ - 1) + // BLOCK_SEQ + ) + + sum_exp = 0.0 + max_logic = -float("inf") + acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32) + + offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d + offs_logic = cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh + for block_seq_n in range(0, block_n_size, 1): + tv = tl.load(Mid_O + offs_v + block_seq_n * stride_mid_os) + tlogic = tl.load(Mid_O_LogExpSum + offs_logic + block_seq_n) + new_max_logic = tl.maximum(tlogic, max_logic) + + old_scale = tl.exp(max_logic - new_max_logic) + acc *= old_scale + exp_logic = tl.exp(tlogic - new_max_logic) + acc += exp_logic * tv + sum_exp = sum_exp * old_scale + exp_logic + max_logic = new_max_logic + + tl.store(O + cur_batch * stride_obs + cur_head * stride_oh + offs_d, acc / sum_exp) + return + + +@torch.no_grad() +def flash_decode_stage1( + q, + k, + v, + Req_to_tokens, + B_req_idx, + B_Seqlen, + max_len_in_batch, + mid_out, + mid_out_logsumexp, + block_seq, +): + BLOCK_SEQ = block_seq + BLOCK_N = 16 + assert BLOCK_SEQ % BLOCK_N == 0 + # shape constraints + Lq, Lk = q.shape[-1], k.shape[-1] + assert Lq == Lk + assert Lk in {16, 32, 64, 128} + sm_scale = 1.0 / (Lk**0.5) + batch, head_num = B_req_idx.shape[0], q.shape[1] + grid = (batch, head_num, triton.cdiv(max_len_in_batch, BLOCK_SEQ)) + gqa_group_size = q.shape[1] // k.shape[1] + + _fwd_kernel_flash_decode_stage1[grid]( + q, + k, + v, + sm_scale, + Req_to_tokens, + B_req_idx, + B_Seqlen, + mid_out, + mid_out_logsumexp, + Req_to_tokens.stride(0), + Req_to_tokens.stride(1), + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + mid_out.stride(0), + mid_out.stride(1), + mid_out.stride(2), + mid_out.stride(3), + mid_out_logsumexp.stride(0), + mid_out_logsumexp.stride(1), + mid_out_logsumexp.stride(2), + gqa_group_size, + BLOCK_SEQ=BLOCK_SEQ, + BLOCK_DMODEL=Lk, + BLOCK_N=BLOCK_N, + num_warps=1, + num_stages=2, + ) + return + + +@torch.no_grad() +def flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, O, block_seq): + Lk = mid_out.shape[-1] + assert Lk in {16, 32, 64, 128} + batch, head_num = mid_out.shape[0], mid_out.shape[1] + grid = (batch, head_num) + + _fwd_kernel_flash_decode_stage2[grid]( + B_Seqlen, + mid_out, + mid_out_logexpsum, + O, + mid_out.stride(0), + mid_out.stride(1), + mid_out.stride(2), + mid_out.stride(3), + mid_out_logexpsum.stride(0), + mid_out_logexpsum.stride(1), + mid_out_logexpsum.stride(2), + O.stride(0), + O.stride(1), + O.stride(2), + BLOCK_SEQ=block_seq, + BLOCK_DMODEL=Lk, + num_warps=4, + num_stages=2, + ) + return + + +def flash_decode_attention_fwd( + q, + k_buffer, + v_buffer, + o, + req_to_token, + b_req_idx, + b_start_loc, + b_seq_len, + attn_logits, + max_len_in_batch, + sm_scale, + logit_cap=0.0, +): + BLOCK_SEQ = 256 + kv_group_num = q.shape[1] // v_buffer.shape[1] + # batch_size = q.shape[0] + + block_seq_num = (max_len_in_batch + BLOCK_SEQ - 1) // BLOCK_SEQ + + mid_o = torch.empty( + [q.shape[0], q.shape[1], block_seq_num, q.shape[-1]], + dtype=torch.float32, + device="cuda", + ) + mid_o_logexpsum = torch.empty( + [q.shape[0], q.shape[1], block_seq_num], dtype=torch.float32, device="cuda" + ) + + flash_decode_stage1( + q, + k_buffer, + v_buffer, + req_to_token, + b_req_idx, + b_seq_len, + max_len_in_batch, + mid_o, + mid_o_logexpsum, + BLOCK_SEQ, + ) + flash_decode_stage2(mid_o, mid_o_logexpsum, b_seq_len, o, BLOCK_SEQ) + + +@triton.jit +def _sparse_fwd_kernel_flash_decode_stage1( # Double Sparsity's approximate attention + Q_Label, + K_Label_Buffer, + sm_scale, + Req_to_tokens, # shape: [B, S] + B_Seqlen, + Att_Out, # shape: [H, B, S] easier for topk + stride_req_to_tokens_b, + stride_qbs, + stride_qh, + stride_buf_kbs, + stride_buf_kh, + att_stride_h, + att_stride_b, + kv_group_num: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + logit_cap: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + start_n = tl.program_id(2) + + cur_kv_head = cur_head // kv_group_num + + offs_d = tl.arange(0, BLOCK_DMODEL) + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + + cur_batch_start_index = 0 + cur_batch_end_index = cur_batch_seq_len + + min_val = -float("inf") + att_value = tl.full([BLOCK_N], min_val, dtype=tl.float32) + + off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d + + offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N) + + block_index = start_n * BLOCK_N + block_mask = tl.where(block_index < cur_batch_seq_len, 1, 0) + + for start_mark in range(0, block_mask, 1): + q = tl.load(Q_Label + off_q + start_mark).to(REDUCE_TRITON_TYPE) + offs_n_new = cur_batch_start_index + offs_n + k_loc = tl.load( + Req_to_tokens + stride_req_to_tokens_b * cur_batch + offs_n_new, + mask=offs_n_new < cur_batch_end_index, + other=0, + ) + offs_buf_k = ( + k_loc[:, None] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_d[None, :] + ) + k = tl.load( + K_Label_Buffer + offs_buf_k, + mask=offs_n_new[:, None] < cur_batch_end_index, + other=0.0, + ).to(REDUCE_TRITON_TYPE) + + att_value = tl.sum(q[None, :] * k, 1) + att_value *= sm_scale + + if logit_cap > 0: + att_value = logit_cap * tanh(att_value / logit_cap) + + att_value = tl.where(offs_n < cur_batch_end_index, att_value, min_val) + off_o = cur_head * att_stride_h + (cur_batch * att_stride_b + offs_n) + tl.store(Att_Out + off_o, att_value) + + +@triton.jit +def _sparse_fwd_kernel_flash_decode_stage2( + Q, + K, + V, + sm_scale, + Req_to_tokens, # shape: [B, S] + Topk_token_indices, # shape: [H, B, k] + Mid_O, # [batch, head, seq_block_num, head_dim] + Mid_O_LogExpSum, # [batch, head, seq_block_num] + Heavy_token_num, # NOTE: This can be used as constexpr but we may support dynamic heavy token number in the future + stride_req_to_tokens_b, + stride_topk_token_indices_h, + stride_topk_token_indices_b, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_mid_o_eb, + stride_mid_o_eh, + gqa_group_size, + BLOCK_SEQ: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + seq_start_block = tl.program_id(2) + cur_kv_head = cur_head // gqa_group_size + + offs_d = tl.arange(0, BLOCK_DMODEL) + cur_batch_start_index = seq_start_block * BLOCK_SEQ + cur_batch_end_index = tl.minimum(Heavy_token_num, cur_batch_start_index + BLOCK_SEQ) + + off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d + + block_n_size = ( + tl.where( + cur_batch_end_index - cur_batch_start_index <= 0, + 0, + cur_batch_end_index - cur_batch_start_index + BLOCK_N - 1, + ) + // BLOCK_N + ) + + # offs_n = cur_batch_start_index + tl.arange(0, BLOCK_N) + offs_n = tl.arange(0, BLOCK_N) + + q = tl.load(Q + off_q) + + sum_exp = 0.0 + max_logic = -float("inf") + acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32) + + for start_n in range(cur_batch_start_index, cur_batch_end_index, BLOCK_N): + # for start_n in range(0, block_n_size, 1): + # offs_n_new = start_n * BLOCK_N + offs_n + offs_n_new = start_n + offs_n + # offs_n_new = cur_batch_start_index + start_n * BLOCK_N + offs_n + topk_token_indices = tl.load( + Topk_token_indices + + stride_topk_token_indices_h * cur_head + + stride_topk_token_indices_b * cur_batch + + offs_n_new, + mask=offs_n_new < cur_batch_end_index, + other=0, + ) + k_loc = tl.load( + Req_to_tokens + stride_req_to_tokens_b * cur_batch + topk_token_indices, + mask=offs_n_new < cur_batch_end_index, + other=0, + ) + off_k = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :] + k = tl.load( + K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0 + ) + att_value = tl.sum(q[None, :] * k, 1) + att_value *= sm_scale + att_value = tl.where(offs_n_new < cur_batch_end_index, att_value, float("-inf")) + v = tl.load( + V + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0 + ) + + cur_max_logic = tl.max(att_value, axis=0) + new_max_logic = tl.maximum(cur_max_logic, max_logic) + + exp_logic = tl.exp(att_value - new_max_logic) + logic_scale = tl.exp(max_logic - new_max_logic) + acc *= logic_scale + acc += tl.sum(exp_logic[:, None] * v, axis=0) + + sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=0) + max_logic = new_max_logic + + # need_store = tl.where(block_n_size == 0, 0, 1) + need_store = 1 + for _ in range(0, need_store, 1): + off_mid_o = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + seq_start_block * stride_mid_os + + offs_d + ) + off_mid_o_logexpsum = ( + cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh + seq_start_block + ) + tl.store(Mid_O + off_mid_o, acc / sum_exp) + tl.store(Mid_O_LogExpSum + off_mid_o_logexpsum, max_logic + tl.log(sum_exp)) + return + + +@triton.jit +def _sparse_fwd_kernel_flash_decode_stage3( + Mid_O, # [batch, head, seq_block_num, head_dim] + Mid_O_LogExpSum, # [batch, head, seq_block_num] + O, # [batch, head, head_dim] + seq_len, # NOTE: This can be used as constexpr but we may support dynamic heavy token number in the future + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_mid_o_eb, + stride_mid_o_eh, + stride_obs, + stride_oh, + BLOCK_SEQ: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + + offs_d = tl.arange(0, BLOCK_DMODEL) + + block_n_size = tl.where(seq_len <= 0, 0, seq_len + BLOCK_SEQ - 1) // BLOCK_SEQ + + sum_exp = 0.0 + max_logic = -float("inf") + acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32) + + offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d + offs_logic = cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh + for block_seq_n in range(0, block_n_size, 1): + tv = tl.load(Mid_O + offs_v + block_seq_n * stride_mid_os) + tlogic = tl.load(Mid_O_LogExpSum + offs_logic + block_seq_n) + new_max_logic = tl.maximum(tlogic, max_logic) + + old_scale = tl.exp(max_logic - new_max_logic) + acc *= old_scale + exp_logic = tl.exp(tlogic - new_max_logic) + acc += exp_logic * tv + sum_exp = sum_exp * old_scale + exp_logic + max_logic = new_max_logic + + tl.store(O + cur_batch * stride_obs + cur_head * stride_oh + offs_d, acc / sum_exp) + return + + +def sparse_flash_decode_stage1( + q_label, + k_label_buffer, + att_out, + Req_to_tokens, + B_Seqlen, + max_len_in_batch, + sm_scale, + logit_cap, +): + BLOCK = 32 + # shape constraints + Lq, Lk = q_label.shape[-1], k_label_buffer.shape[-1] + assert Lq == Lk + assert Lk in {16, 32, 64, 128, 256, 576} + + BLOCK_DMODEL = Lk + + batch, head_num = q_label.shape[0], q_label.shape[1] + + grid = (batch, head_num, triton.cdiv(max_len_in_batch, BLOCK)) + kv_group_num = q_label.shape[1] // k_label_buffer.shape[1] + + if kv_group_num == 1: + num_warps = 4 + else: + num_warps = 2 + + _sparse_fwd_kernel_flash_decode_stage1[grid]( + q_label, + k_label_buffer, + sm_scale, + Req_to_tokens, + B_Seqlen, + att_out, + Req_to_tokens.stride(0), + q_label.stride(0), + q_label.stride(1), + k_label_buffer.stride(0), + k_label_buffer.stride(1), + att_out.stride(0), + att_out.stride(1), + kv_group_num, + BLOCK_DMODEL, + BLOCK, + logit_cap, + num_warps=num_warps, + num_stages=1, + ) + + +@torch.no_grad() +def sparse_flash_decode_stage2( + q, + k, + v, + Req_to_tokens, + Topk_token_indices, + heavy_token_num, + mid_out, + mid_out_logsumexp, + block_seq, + sm_scale, +): + BLOCK_SEQ = block_seq + BLOCK_N = 16 + assert BLOCK_SEQ % BLOCK_N == 0 + # shape constraints + Lq, Lk = q.shape[-1], k.shape[-1] + assert Lq == Lk + assert Lk in {16, 32, 64, 128} + assert heavy_token_num == Topk_token_indices.shape[-1] + # sm_scale = 1.0 / (Lk ** 0.5) + batch, head_num = q.shape[0], q.shape[1] + grid = (batch, head_num, triton.cdiv(heavy_token_num, BLOCK_SEQ)) + + gqa_group_size = q.shape[1] // k.shape[1] + + _sparse_fwd_kernel_flash_decode_stage2[grid]( + q, + k, + v, + sm_scale, + Req_to_tokens, + Topk_token_indices, + mid_out, + mid_out_logsumexp, + heavy_token_num, + Req_to_tokens.stride(0), + Topk_token_indices.stride(0), + Topk_token_indices.stride(1), + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + mid_out.stride(0), + mid_out.stride(1), + mid_out.stride(2), + mid_out_logsumexp.stride(0), + mid_out_logsumexp.stride(1), + gqa_group_size, + BLOCK_SEQ=BLOCK_SEQ, + BLOCK_DMODEL=Lk, + BLOCK_N=BLOCK_N, + num_warps=1, + num_stages=2, + ) + return + + +@torch.no_grad() +def sparse_flash_decode_stage3(Seqlen, mid_out, mid_out_logexpsum, O, block_seq): + Lk = mid_out.shape[-1] + assert Lk in {16, 32, 64, 128} + batch, head_num = mid_out.shape[0], mid_out.shape[1] + grid = (batch, head_num) + + _sparse_fwd_kernel_flash_decode_stage3[grid]( + mid_out, + mid_out_logexpsum, + O, + Seqlen, + mid_out.stride(0), + mid_out.stride(1), + mid_out.stride(2), + mid_out_logexpsum.stride(0), + mid_out_logexpsum.stride(1), + O.stride(0), + O.stride(1), + BLOCK_SEQ=block_seq, + BLOCK_DMODEL=Lk, + num_warps=4, + num_stages=2, + ) + return + + +def flash_decode_sparse_attention_fwd( + q, + k_buffer, + v_buffer, + o, + q_label, + k_label_buffer, + req_to_token, + b_seq_len, + max_len_in_batch, + sm_scale, + logit_cap, + heavy_token_num=32, + att_out_approx=None, + mid_out=None, + mid_o_logexpsum=None, + BLOCK_SEQ=256, +): + # TODO(Andy): Tune BLOCK_SEQ & BLOCK_D + kv_group_num = q.shape[1] // v_buffer.shape[1] + # batch_size = q.shape[0] + + # Step 1: BGEMV approximate attention (page implementation) + + if att_out_approx is None: + att_out_approx = torch.empty( + [q.shape[1], q.shape[0], max_len_in_batch], + dtype=REDUCE_TORCH_TYPE, + device=q.device, + ) + + if mid_out is None: + block_seq_num = (heavy_token_num + BLOCK_SEQ - 1) // BLOCK_SEQ + + mid_out = torch.empty( + [q.shape[0], q.shape[1], block_seq_num, q.shape[-1]], + dtype=torch.float32, + device=q.device, + ) + mid_o_logexpsum = torch.empty( + [q.shape[0], q.shape[1], block_seq_num], + dtype=torch.float32, + device=q.device, + ) + + sparse_flash_decode_stage1( + q_label, + k_label_buffer, + att_out_approx, + req_to_token, + b_seq_len, + max_len_in_batch, + sm_scale, + logit_cap, + ) + + # Step 2: TopK token selection + # NOTE(Andy): Apply sparse decoding when min > heavy_token_num and max > sparse decoding threshold + # TODO(Andy): Change a faster topk implementation + topk_token_indices = torch.topk(att_out_approx, heavy_token_num, dim=-1).indices + # topk_token_indices: [H, B, k], Req_to_tokens: [B, S] + # topk_token_indices = torch.arange(0, heavy_token_num, device=q.device).unsqueeze(0).unsqueeze(0).expand(q.shape[1], q.shape[0], -1) + + sparse_flash_decode_stage2( + q, + k_buffer, + v_buffer, + req_to_token, + topk_token_indices, + heavy_token_num, + mid_out, + mid_o_logexpsum, + BLOCK_SEQ, + sm_scale, + ) + + sparse_flash_decode_stage3(heavy_token_num, mid_out, mid_o_logexpsum, o, BLOCK_SEQ) + + +# Extend attention kernel for Double Sparsity +# Moved from https://github.com/sgl-project/sglang/blob/v0.4.2.post1/python/sglang/srt/layers/attention/triton_ops/extend_attention.py +@triton.jit +def _fwd_kernel( + Q_Extend, + K_Extend, + V_Extend, + O_Extend, + K_Buffer, + V_Buffer, + Req_to_tokens, + B_req_idx, + B_Seq_Len, + B_Start_Loc_Extend, + B_Seq_Len_Extend, + 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, + stride_req_to_tokens_b, + 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, +): + 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_len = tl.load(B_Seq_Len + cur_seq) + cur_seq_len_extend = tl.load(B_Seq_Len_Extend + cur_seq) + cur_seq_len_prefix = cur_seq_len - cur_seq_len_extend + + cur_seq_prefix_start_in_loc = 0 + cur_seq_extend_start_contiguous = tl.load(B_Start_Loc_Extend + cur_seq) + cur_batch_req_idx = tl.load(B_req_idx + 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_contiguous + 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_contiguous + 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 + offs_b_loc_prefix = cur_batch_req_idx * stride_req_to_tokens_b + ( + cur_seq_prefix_start_in_loc + start_n + offs_n + ) + offs_kv_loc = tl.load(Req_to_tokens + offs_b_loc_prefix, 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(mask_m[:, None] & mask_n[None, :], qk, float("-inf")) + + 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]) + 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 = 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 + + # load k in transposed way + offs_k = ( + (cur_seq_extend_start_contiguous + 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_contiguous + 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) + + mask_causual = (cur_block_m * BLOCK_M + offs_m[:, None]) >= ( + start_n + offs_n[None, :] + ) + mask_causual &= mask_m[:, None] & mask_n[None, :] + qk = tl.where(mask_causual, qk, float("-inf")) + + 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]) + deno = deno * re_scale + tl.sum(p, 1) + + offs_v = ( + (cur_seq_extend_start_contiguous + 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 + + offs_o = ( + (cur_seq_extend_start_contiguous + cur_block_m * BLOCK_M + offs_m[:, None]) + * stride_obs + + cur_head * stride_oh + + offs_dv[None, :] + ) + 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, + req_to_tokens, + b_req_idx, + b_seq_len, + b_seq_len_extend, + b_start_loc_extend, + max_len_extend, + sm_scale=None, + logit_cap=0.0, +): + """ + 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: + 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 = b_seq_len.shape[0], q_extend.shape[1] + kv_group_num = q_extend.shape[1] // k_extend.shape[1] + + 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": 4, "matrix_instr_nonkdim": 16, "kpack": 2} + + _fwd_kernel[grid]( + q_extend, + k_extend, + v_extend, + o_extend, + k_buffer, + v_buffer, + req_to_tokens, + b_req_idx, + b_seq_len, + b_start_loc_extend, + b_seq_len_extend, + 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), + req_to_tokens.stride(0), + 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, + num_warps=num_warps, + num_stages=num_stages, + **extra_kargs, + ) diff --git a/python/sglang/srt/layers/attention/triton_ops/extend_attention.py b/python/sglang/srt/layers/attention/triton_ops/extend_attention.py new file mode 100644 index 000000000..e91467743 --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/extend_attention.py @@ -0,0 +1,550 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +""" +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 + +from sglang.srt.layers.attention.triton_ops.prefill_attention import ( + context_attention_fwd, +) +from sglang.srt.utils import is_cuda, is_hip + +_is_cuda = is_cuda() +if _is_cuda: + CUDA_CAPABILITY = torch.cuda.get_device_capability() + +_is_hip = is_hip() + + +@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, + window_kv_offset_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, + xai_temperature_len: 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) + + # For SWA, we should only load the mask in the sliding window + window_kv_offset = 0 + if USE_CUSTOM_MASK and SLIDING_WINDOW_SIZE > 0: + window_kv_offset = tl.load(window_kv_offset_ptr + 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 + + if xai_temperature_len > 0: + offs_qidx = cur_seq_len_prefix + cur_block_m * BLOCK_M + offs_m + xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len)) + xai_temperature_reg = tl.where( + offs_qidx > xai_temperature_len, + tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale, + 1.0, + ) + + 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 + window_kv_offset) + + window_kv_offset + + 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) + + if xai_temperature_len > 0: + qk *= xai_temperature_reg[:, None] + + 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 + window_kv_offset) + + window_kv_offset + + 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) + + if xai_temperature_len > 0: + qk *= xai_temperature_reg[:, None] + + 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, + window_kv_offsets=None, + xai_temperature_len=-1, +): + """ + 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, + window_kv_offsets, + 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, + xai_temperature_len=xai_temperature_len, + 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 redundant_attention( + q_extend, + o_extend, + k_buffer, + v_buffer, + b_req_idx, + b_start_loc, + b_seq_len, + b_seq_len_prefix, + max_len_in_batch, +): + total_token_num = k_buffer.shape[0] + B, H_Q, D = b_req_idx.shape[0], q_extend.shape[-2], q_extend.shape[-1] + q_buffer = torch.empty( + (total_token_num, H_Q, D), dtype=q_extend.dtype, device=q_extend.device + ) + + pt = 0 + for i in range(B): + cur_seq_len_extend = b_seq_len[i] - b_seq_len_prefix[i] + pl, pr = b_start_loc[i] + b_seq_len_prefix[i], b_start_loc[i] + b_seq_len[i] + q_buffer[pl:pr] = q_extend[pt : pt + cur_seq_len_extend] + pt += cur_seq_len_extend + + o_buffer = torch.empty_like(q_buffer) + context_attention_fwd( + q_buffer, k_buffer, v_buffer, o_buffer, b_start_loc, b_seq_len, max_len_in_batch + ) + + pt = 0 + for i in range(B): + cur_seq_len_extend = b_seq_len[i] - b_seq_len_prefix[i] + pl, pr = b_start_loc[i] + b_seq_len_prefix[i], b_start_loc[i] + b_seq_len[i] + o_extend[pt : pt + cur_seq_len_extend] = o_buffer[pl:pr] + pt += cur_seq_len_extend diff --git a/python/sglang/srt/layers/attention/triton_ops/merge_state.py b/python/sglang/srt/layers/attention/triton_ops/merge_state.py new file mode 100644 index 000000000..955097fc9 --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/merge_state.py @@ -0,0 +1,96 @@ +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + + +@triton.jit +def merge_state_kernel( + output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_merged + output_lse, # [NUM_TOKENS, NUM_HEADS] s_merged + prefix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_a + prefix_lse, # [NUM_TOKENS, NUM_HEADS] s_a + suffix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_b + suffix_lse, # [NUM_TOKENS, NUM_HEADS] s_b + HEAD_SIZE: tl.constexpr, + PADDED_HEAD_SIZE: tl.constexpr, + OUTPUT_LSE: tl.constexpr, +): + token_idx = tl.program_id(0) + num_tokens = tl.num_programs(0) + head_idx = tl.program_id(1) + num_heads = tl.num_programs(1) + + p_lse = tl.load(prefix_lse + token_idx * num_heads + head_idx) + s_lse = tl.load(suffix_lse + token_idx * num_heads + head_idx) + p_lse = float("-inf") if p_lse == float("inf") else p_lse + s_lse = float("-inf") if s_lse == float("inf") else s_lse + + max_lse = tl.maximum(p_lse, s_lse) + p_lse = p_lse - max_lse + s_lse = s_lse - max_lse + out_se = tl.exp(p_lse) + tl.exp(s_lse) + + if OUTPUT_LSE: + out_lse = tl.log(out_se) + max_lse + tl.store(output_lse + token_idx * num_heads + head_idx, out_lse) + + head_arange = tl.arange(0, PADDED_HEAD_SIZE) + head_mask = head_arange < HEAD_SIZE + p_out = tl.load( + prefix_output + + token_idx * num_heads * HEAD_SIZE + + head_idx * HEAD_SIZE + + head_arange, + mask=head_mask, + ) + s_out = tl.load( + suffix_output + + token_idx * num_heads * HEAD_SIZE + + head_idx * HEAD_SIZE + + head_arange, + mask=head_mask, + ) + + p_scale = tl.exp(p_lse) / out_se + s_scale = tl.exp(s_lse) / out_se + out = p_out * p_scale + s_out * s_scale + tl.store( + output + token_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE + head_arange, + out, + mask=head_mask, + ) + + +def merge_state_triton( + prefix_output: torch.Tensor, + prefix_lse: torch.Tensor, + suffix_output: torch.Tensor, + suffix_lse: torch.Tensor, + output: Optional[torch.Tensor] = None, + output_lse: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + # Avoid creating new tensors if they are already provided + if output is None: + output = torch.empty_like(prefix_output) + if output_lse is None: + output_lse = torch.empty_like(prefix_lse) + + num_tokens = output.shape[0] + num_query_heads = output.shape[1] + head_size = output.shape[2] + padded_head_size = triton.next_power_of_2(head_size) + + merge_state_kernel[(num_tokens, num_query_heads)]( + output, + output_lse, + prefix_output, + prefix_lse, + suffix_output, + suffix_lse, + head_size, + padded_head_size, + output_lse is not None, + ) + return output, output_lse diff --git a/python/sglang/srt/layers/attention/triton_ops/prefill_attention.py b/python/sglang/srt/layers/attention/triton_ops/prefill_attention.py new file mode 100644 index 000000000..ac0fc72af --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/prefill_attention.py @@ -0,0 +1,217 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +""" +Memory-efficient attention for prefill. +It supporst page size = 1. +""" + +# Adapted from +# https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py#L1 +import torch +import triton +import triton.language as tl + +from sglang.srt.utils import is_cuda, is_hip + +_is_cuda = is_cuda() +_is_hip = is_hip() + +if _is_cuda or _is_hip: + CUDA_CAPABILITY = torch.cuda.get_device_capability() + + +@triton.jit +def _fwd_kernel( + Q, + K, + V, + sm_scale, + B_Start_Loc, + B_Seqlen, + Out, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_obs, + stride_oh, + kv_group_num: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, + Lk: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + start_m = tl.program_id(2) + + cur_kv_head = cur_head // kv_group_num + + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) + + block_start_loc = BLOCK_M * start_m + + # initialize offsets + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + off_q = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] + ) + off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] + off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] + + mask_d = offs_d < Lk + + q = tl.load( + Q + off_q, + mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]), + other=0.0, + ) + + k_ptrs = K + off_k + v_ptrs = V + off_v + + # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + + block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0) + + end_n = ( + cur_batch_seq_len + if not IS_CAUSAL + else tl.minimum((start_m + 1) * BLOCK_M, cur_batch_seq_len) + ) + for start_n in range(0, block_mask * end_n, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + k = tl.load( + k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs, + mask=((start_n + offs_n[None, :]) < cur_batch_seq_len) & (mask_d[:, None]), + other=0.0, + ) + # mask = tl.load(mask_ptrs + start_n, mask=start_n + offs_n < cur_batch_end_loc, other=0.0) + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk += tl.dot(q, k) + qk *= sm_scale + + if IS_CAUSAL: + qk += tl.where( + (start_n + offs_n[None, :] < cur_batch_seq_len) + & (offs_m[:, None] >= (start_n + offs_n[None, :])), + 0, + float("-inf"), + ) + else: + qk += tl.where( + (start_n + offs_n[None, :]) < cur_batch_seq_len, 0, float("-inf") + ) + + # -- compute m_ij, p, l_ij + m_ij = tl.max(qk, 1) + p = tl.exp(qk - m_ij[:, None]) + l_ij = tl.sum(p, 1) + # -- update m_i and l_i + m_i_new = tl.maximum(m_i, m_ij) + alpha = tl.exp(m_i - m_i_new) + beta = tl.exp(m_ij - m_i_new) + l_i_new = alpha * l_i + beta * l_ij + # -- update output accumulator -- + # scale p + p_scale = beta / l_i_new + p = p * p_scale[:, None] + # scale acc + acc_scale = l_i / l_i_new * alpha + acc = acc * acc_scale[:, None] + # update acc + v = tl.load( + v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs, + mask=((start_n + offs_n[:, None]) < cur_batch_seq_len) & (mask_d[None, :]), + other=0.0, + ) + + p = p.to(v.dtype) + acc += tl.dot(p, v) + # update m_i and l_i + l_i = l_i_new + m_i = m_i_new + # initialize pointers to output + off_o = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + + cur_head * stride_oh + + offs_d[None, :] + ) + out_ptrs = Out + off_o + tl.store( + out_ptrs, acc, mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]) + ) + + +def context_attention_fwd( + q, k, v, o, b_start_loc, b_seq_len, max_input_len, is_causal=True +): + """ + q, k, v: [b * s, head, head_dim] + b_start_loc: [b] + b_seq_len: [b] + out: [b * s, head, head_dim] + """ + if (_is_cuda or _is_hip) and CUDA_CAPABILITY[0] > 8: + BLOCK = 128 + else: + BLOCK = 64 + + Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] + + sm_scale = 1.0 / (Lq**0.5) + batch, head = b_seq_len.shape[0], q.shape[1] + kv_group_num = q.shape[1] // k.shape[1] + + grid = (batch, head, triton.cdiv(max_input_len, BLOCK)) + num_warps = 4 if Lk <= 64 else 8 + + _fwd_kernel[grid]( + q, + k, + v, + sm_scale, + b_start_loc, + b_seq_len, + o, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + o.stride(0), + o.stride(1), + kv_group_num=kv_group_num, + BLOCK_M=BLOCK, + BLOCK_DMODEL=triton.next_power_of_2(Lk), + BLOCK_N=BLOCK, + IS_CAUSAL=is_causal, + num_warps=num_warps, + num_stages=1, + Lk=Lk, + ) diff --git a/python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py b/python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py new file mode 100644 index 000000000..89b17ce2d --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py @@ -0,0 +1,439 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +""" +Memory-efficient attention for decoding. +It supports page size = 1. +""" + +# Adapted from +# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py +# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py + +import triton +import triton.language as tl + +from sglang.srt.layers.attention.triton_ops.decode_attention import ( + _decode_softmax_reducev_fwd, +) + + +def is_hip(): + return triton.runtime.driver.active.get_current_target().backend == "hip" + + +_is_hip = is_hip() + + +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def _fwd_grouped_kernel_stage1_rope( + Q, # Holds [Q_NOPE; Q_PE], b x h x (d+r) + K_Buffer, # Holds [KV; K_PE], b*s x (c+r) + V_buffer, # Holds [KV], b*s x (c) + cos_sin_cache, # max_seq_len x (rotary_dim * 2) + positions, # sequence positions + sm_scale, + kv_indptr, + kv_indices, + Att_Out, # b x h x NUM_KV_SPLITS x (kv_lora_rank + 1) + k_pe_t_out, + stride_qb, + stride_qh, + stride_buf_kbs, + stride_buf_vbs, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_kpe_tokens_out_b, + stride_cos_sin_cache_s, + stride_positions_b, + rotary_dim: tl.constexpr, + kv_lora_rank: tl.constexpr, + qk_rope_head_dim: tl.constexpr, + kv_group_num: tl.constexpr, + q_head_num: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_R: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, + logit_cap: tl.constexpr, + USE_ROPE: tl.constexpr, + IS_NEOX_STYLE: tl.constexpr, +): + + cur_batch = tl.program_id(0) + cur_head_id = tl.program_id(1) + split_kv_id = tl.program_id(2) + + 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_c = tl.arange(0, BLOCK_C) + offs_qk_r = tl.arange(kv_lora_rank, kv_lora_rank + BLOCK_R) # to get the k_pe + + off_q_pe = ( + cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_qk_r[None, :] + ) + offs_q = cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_c[None, :] + + mask_c = offs_c < kv_lora_rank + mask_qk_r = offs_qk_r < (kv_lora_rank + qk_rope_head_dim) + + 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 + + q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_c[None, :]), other=0.0) + q_pe = tl.load( + Q + off_q_pe, mask=(mask_h[:, None]) & (mask_qk_r[None, :]), other=0.0 + ) + + kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_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) + + # apply rotary embedding for q_pe, and k_pe (last token per batch of K_PE) + LAST_SPLIT = split_kv_end == cur_batch_seq_len + k_pe_last_token = tl.zeros([BLOCK_R], dtype=q.dtype) + + if USE_ROPE: + if IS_NEOX_STYLE: + # [BLOCK_ROTARY // 2, BLOCK_ROTARY // 2 + 1, BLOCK_ROTARY // 2 + 2, ..., 0, 1, 2, ..., BLOCK_ROTARY // 2 - 1, pass:] + offs_qk_rot_r = kv_lora_rank + ( + (tl.arange(0, BLOCK_R) + (rotary_dim // 2)) % rotary_dim + ) + # Which elements to flip + mask_rotate = tl.arange(0, BLOCK_R) < (rotary_dim // 2) + # [0 , 1, 2, ..., rotary_dim // 2 - 1, 0 , 1, 2, ..., rotary_dim // 2 - 1] + offs_rotary = tl.arange(0, BLOCK_R) % (rotary_dim // 2) + else: + # [1, 0, 3, 2, 5, 4, ..., BLOCK_R, BLOCK_R - 1] + offs_qk_rot_r = ( + kv_lora_rank + + (((tl.arange(0, BLOCK_R) + 1) % 2) * 2) + - 1 + + tl.arange(0, BLOCK_R) + ) + mask_rotate = tl.arange(0, BLOCK_R) % 2 < 1 + # [0, 0, 1, 1, ..., rotary_dim // 2 - 1, rotary_dim // 2 - 1] + offs_rotary = tl.arange(0, BLOCK_R) // 2 + + if qk_rope_head_dim > rotary_dim: + offs_qk_rot_r = tl.where( + tl.arange(0, BLOCK_R) < rotary_dim, offs_qk_rot_r, tl.arange(0, BLOCK_R) + ) + offs_rotary = tl.where( + tl.arange(0, BLOCK_R) < rotary_dim, offs_rotary, tl.arange(0, BLOCK_R) + ) + + mask_rotary = tl.arange(0, BLOCK_R) < rotary_dim + + pos = tl.load(positions + cur_batch * stride_positions_b) + cos = tl.load( + cos_sin_cache + pos * stride_cos_sin_cache_s + offs_rotary, + mask=mask_rotary, + other=1.0, + ) + sin = tl.load( + cos_sin_cache + + pos * stride_cos_sin_cache_s + + offs_rotary + + rotary_dim // 2, + mask_rotary, + other=0.0, + ) + + off_q_pe_rot = ( + cur_batch * stride_qb + + cur_head[:, None] * stride_qh + + offs_qk_rot_r[None, :] + ) + mask_qk_rot_r = offs_qk_rot_r < (kv_lora_rank + qk_rope_head_dim) + + # 0, 2, 4,.... 1, 3, 5... + q_pe_rot = tl.load( + Q + off_q_pe_rot, + mask=(mask_h[:, None]) & (mask_qk_rot_r[None, :]), + other=0.0, + ) + q_pe_rot = tl.where(mask_rotate[None, :], -q_pe_rot, q_pe_rot) + + q_pe = q_pe * cos + q_pe_rot * sin + + # we only apply to the last token in the K_PE + if LAST_SPLIT: + # debug assert + if (cur_batch == 0 and cur_head == 0) and split_kv_id < NUM_KV_SPLITS - 1: + tl.device_assert(False, "Only last split should compute k_pe") + + kv_loc = tl.load( + kv_indices + cur_batch_kv_start_idx + cur_batch_seq_len - 1 + ) + offs_buf_k_pe_last_token = kv_loc * stride_buf_kbs + offs_qk_r + offs_buf_k_pe_rot_last_token = kv_loc * stride_buf_kbs + offs_qk_rot_r + k_pe_last_token = tl.load(K_Buffer + offs_buf_k_pe_last_token) + + k_pe_rot_last_token = tl.load(K_Buffer + offs_buf_k_pe_rot_last_token) + k_pe_rot_last_token = tl.where( + mask_rotate, -k_pe_rot_last_token, k_pe_rot_last_token + ) + + k_pe_last_token = k_pe_last_token * cos + k_pe_rot_last_token * sin + + 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_C], dtype=tl.float32) + + if split_kv_end > split_kv_start: + 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_kv = kv_loc[None, :] * stride_buf_kbs + offs_c[:, None] + offs_buf_k_pe = kv_loc[None, :] * stride_buf_kbs + offs_qk_r[:, None] + + k_pe = tl.load( + K_Buffer + offs_buf_k_pe, + mask=(offs_n[None, :] < split_kv_end) & (mask_qk_r[:, None]), + other=0.0, + ) # positional embedding part of keys + + if (USE_ROPE and LAST_SPLIT) and start_n >= cur_batch_seq_len - BLOCK_N: + k_pe = tl.where( + offs_n[None, :] != (split_kv_end - 1), + k_pe, + k_pe_last_token[:, None], + ) + + # (16, 64) x (64, 32) + # dot product of rope parts + qk = tl.dot(q_pe, k_pe.to(q_pe.dtype)) + + kv = tl.load( + K_Buffer + offs_buf_kv, + mask=(offs_n[None, :] < split_kv_end) & (mask_c[:, None]), + other=0.0, + ) # the shared latent tensor for keys and values + + # (16, 512) x (512, 32) + # dot product of nope parts + qk += tl.dot(q, kv) + + 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 + offs_c[None, :] + v = tl.load( + V_buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_c[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] + # (16, 32) x (32, 512) + 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_c[None, :] + ) + + if USE_ROPE: + if LAST_SPLIT: + k_pe_last_token_ptrs = ( + k_pe_t_out + + cur_batch * stride_kpe_tokens_out_b + + tl.arange(0, BLOCK_R) + ) + tl.store(k_pe_last_token_ptrs, k_pe_last_token, mask=mask_qk_r) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum[:, None], + mask=(mask_h[:, None]) & (mask_c[None, :]), + ) + + offs_mid_o_1 = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + + kv_lora_rank + ) + + tl.store( + Att_Out + offs_mid_o_1, + e_max + tl.log(e_sum), + mask=mask_h, + ) + + +# TODO rope offset +def _decode_grouped_att_m_fwd_rope( + q, + k_buffer, + v_buffer, + att_out, + k_pe_tokens_out, + kv_lora_rank, # c + cos_sin_cache, + positions, + rotary_dim, + kv_indptr, + kv_indices, + num_kv_splits, + sm_scale, + logit_cap, + use_rope, + is_neox_style=True, +): + if use_rope: + assert ( + k_pe_tokens_out is not None + ), "We must output the k_pe tokens with rope applied if rope fusion enabled." + + BLOCK = 32 + + # # [TODO] work around shmem limit on MI3xx + # if _is_hip and kv_lora_rank >= 576: + # BLOCK = 16 + + qk_rope_head_dim = k_buffer.shape[-1] - kv_lora_rank + batch, head_num = kv_indptr.shape[0] - 1, q.shape[1] + kv_group_num = q.shape[1] // k_buffer.shape[1] + + BLOCK_C = triton.next_power_of_2(kv_lora_rank) + BLOCK_R = triton.next_power_of_2(qk_rope_head_dim) + + BLOCK_H = 16 + NUM_KV_SPLITS = num_kv_splits + grid = ( + batch, + triton.cdiv(head_num, min(BLOCK_H, kv_group_num)), + NUM_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_rope[grid]( + q, + k_buffer, + v_buffer, + cos_sin_cache, + positions, + sm_scale, + kv_indptr, + kv_indices, + att_out, + k_pe_tokens_out, + q.stride(0), + q.stride(1), + k_buffer.stride(0), + v_buffer.stride(0), + att_out.stride(0), + att_out.stride(1), + att_out.stride(2), + k_pe_tokens_out.stride(0) if use_rope else 0, + cos_sin_cache.stride(0) if use_rope else 0, + positions.stride(0) if use_rope else 0, + rotary_dim, + kv_lora_rank, + qk_rope_head_dim, + kv_group_num=kv_group_num, + q_head_num=head_num, + BLOCK_C=BLOCK_C, + BLOCK_R=BLOCK_R, + BLOCK_N=BLOCK, + BLOCK_H=BLOCK_H, + NUM_KV_SPLITS=NUM_KV_SPLITS, + logit_cap=logit_cap, + USE_ROPE=use_rope, + IS_NEOX_STYLE=is_neox_style, + num_warps=4, + num_stages=num_stages, + **extra_kargs + ) + + +def decode_attention_fwd_grouped_rope( + q, + k_buffer, + v_buffer, + o, + kv_indptr, + kv_indices, + k_pe_tokens, + kv_lora_rank, + rotary_dim, + cos_sin_cache, + positions, + attn_logits, + num_kv_splits, + sm_scale, + logit_cap=0.0, + use_rope=False, + is_neox_style=False, +): + _decode_grouped_att_m_fwd_rope( + q, + k_buffer, + v_buffer, + attn_logits, + k_pe_tokens, + kv_lora_rank, + cos_sin_cache, + positions, + rotary_dim, + kv_indptr, + kv_indices, + num_kv_splits, + sm_scale, + logit_cap, + use_rope, + is_neox_style, + ) + _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, kv_indptr, num_kv_splits) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py new file mode 100644 index 000000000..a48cc9794 --- /dev/null +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -0,0 +1,694 @@ +from __future__ import annotations + +""" +Support attention backend for TRTLLM MHA kernels from flashinfer. +The kernel supports sm100 only, with sliding window and attention sink features. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.attention.flashinfer_backend import ( + FlashInferAttnBackend, + FlashInferMultiStepDraftBackend, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils import is_flashinfer_available + +if is_flashinfer_available(): + import flashinfer + +from sglang.srt.speculative.eagle_utils import EagleDraftInput + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpecInfo + +# Constants +DEFAULT_WORKSPACE_SIZE_MB = ( + 512 # Memory workspace size in MB, todo(Yingyi): read from config +) + +# Reuse this workspace buffer across all TRTLLM MHA wrappers +global_zero_init_workspace_buffer = None + + +@dataclass +class TRTLLMMHAMetadata: + # Sequence lengths for the forward batch + cache_seqlens_int32: torch.Tensor = None + # Maximum sequence length for query + max_seq_len_q: int = 1 + # Maximum sequence length for key + max_seq_len_k: int = 0 + # Cumulative sequence lengths for `query + cu_seqlens_q: torch.Tensor = None + # Cumulative sequence lengths for key + cu_seqlens_k: torch.Tensor = None + # Page table, the index of KV Cache Tables/Blocks + page_table: torch.Tensor = None + + +class TRTLLMHAAttnBackend(FlashInferAttnBackend): + """TRTLLM MHA attention kernel from flashinfer.""" + + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + kv_indptr_buf: Optional[torch.Tensor] = None, + kv_last_page_len_buf: Optional[torch.Tensor] = None, + speculative_step_id: int = 0, + ): + super().__init__( + model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf + ) + + config = model_runner.model_config + + # MHA-specific dimensions + self.max_context_len = model_runner.model_config.context_len + self.hidden_size = config.hidden_size + + # Runtime parameters + self.data_type = model_runner.kv_cache_dtype + self.q_data_type = model_runner.dtype + self.page_size = model_runner.page_size + self.req_to_token = model_runner.req_to_token_pool.req_to_token + self.device = model_runner.device + + # Workspace allocation + self.workspace_size = DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024 + # Allocate buffers + global global_zero_init_workspace_buffer + if global_zero_init_workspace_buffer is None: + global_zero_init_workspace_buffer = torch.zeros( + self.workspace_size, + dtype=torch.uint8, + device=model_runner.device, + ) + self.workspace_buffer = global_zero_init_workspace_buffer + + # CUDA graph state + self.decode_cuda_graph_metadata = {} + + # Speculative decoding + # Only support topk <= 1 for now. + self.topk = model_runner.server_args.speculative_eagle_topk or 0 + self.speculative_step_id = speculative_step_id + self.target_verify_metadata = {} + + self.speculative_num_draft_tokens = ( + model_runner.server_args.speculative_num_draft_tokens + ) + + # Forward metadata + self.forward_metadata: Optional[TRTLLMMHAMetadata] = None + + def init_cuda_graph_state( + self, + max_bs: int, + max_num_tokens: int, + kv_indices_buf: Optional[torch.Tensor] = None, + ): + """Initialize CUDA graph state for TRTLLM MHA.""" + max_num_pages = (self.max_context_len + self.page_size - 1) // self.page_size + self.decode_cuda_graph_metadata = { + "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device), + "page_table": torch.zeros( + max_bs, + max_num_pages, + dtype=torch.int32, + device=self.device, + ), + "strided_indices": torch.arange( + 0, self.max_context_len, self.page_size, device=self.device + ), + } + + if ( + self.speculative_num_draft_tokens is not None + and self.speculative_num_draft_tokens > 0 + ): + self.decode_cuda_graph_metadata["cu_seqlens_q"] = torch.arange( + 0, max_bs + 1, dtype=torch.int32, device=self.device + ) + self.decode_cuda_graph_metadata["cu_seqlens_k"] = torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ) + self.decode_cuda_graph_metadata["page_table_draft_decode"] = torch.zeros( + max_bs, + max_num_pages, + dtype=torch.int32, + device=self.device, + ) + self.target_verify_metadata = { + "cache_seqlens": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "cu_seqlens_q": torch.arange( + 0, + max_bs * self.speculative_num_draft_tokens + 1, + step=self.speculative_num_draft_tokens, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + "page_table": torch.zeros( + max_bs, + max_num_pages, + dtype=torch.int32, + device=self.device, + ), + "strided_indices": torch.arange( + 0, self.max_context_len, self.page_size, device=self.device + ), + } + + self.draft_extend_metadata = { + "cache_seqlens": torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ), + "cu_seqlens_q": torch.zeros( + max_bs + 1, + dtype=torch.int32, + device=self.device, + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + "page_table": torch.zeros( + max_bs, + max_num_pages, + dtype=torch.int32, + device=self.device, + ), + "strided_indices": torch.arange( + 0, self.max_context_len, self.page_size, device=self.device + ), + } + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + ): + """Initialize metadata for CUDA graph capture.""" + metadata = TRTLLMMHAMetadata() + device = seq_lens.device + + if forward_mode.is_decode_or_idle(): + if spec_info is not None: + # Draft Decode + # Here we only support topk = 1 for now. + metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[ + "cache_seqlens" + ][:bs] + metadata.max_seq_len_k = seq_lens.max().item() + ( + self.speculative_step_id + 1 + ) + metadata.cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][ + : bs + 1 + ] + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) + metadata.page_table = self.decode_cuda_graph_metadata[ + "page_table_draft_decode" + ][:bs, :] + self.decode_cuda_graph_metadata[bs] = metadata + else: + # Normal Decode + # Get sequence information + metadata.cache_seqlens_int32 = seq_lens[:bs].to(torch.int32) + batch_size = len(seq_lens) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(seq_lens, dim=0, dtype=torch.int32), (1, 0) + ) + + # Precompute maximum sequence length + metadata.max_seq_len_k = seq_lens.max().item() + # Precompute cumulative sequence lengths + metadata.cu_seqlens_q = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device=device + ) + # Precompute page table + metadata.page_table = self.decode_cuda_graph_metadata["page_table"][ + :bs, : + ] + self.decode_cuda_graph_metadata[bs] = metadata + elif forward_mode.is_target_verify(): + # Target Verify + # Here we only support topk = 1 for now. + metadata.cache_seqlens_int32 = self.target_verify_metadata["cache_seqlens"][ + :bs + ] + metadata.cache_seqlens_int32.copy_( + (seq_lens + self.speculative_num_draft_tokens) + ) + + metadata.cu_seqlens_q = torch.arange( + 0, + bs * self.speculative_num_draft_tokens + 1, + self.speculative_num_draft_tokens, + dtype=torch.int32, + device=device, + ) + + metadata.cu_seqlens_k = self.target_verify_metadata["cu_seqlens_k"][ + : (bs + 1) + ] + + metadata.max_seq_len_q = self.speculative_num_draft_tokens + metadata.max_seq_len_k = ( + seq_lens.max().item() + self.speculative_num_draft_tokens + ) + + metadata.page_table = self.target_verify_metadata["page_table"][:bs, :] + + self.target_verify_metadata[bs] = metadata + elif forward_mode.is_draft_extend(): + metadata.cache_seqlens_int32 = self.draft_extend_metadata["cache_seqlens"][ + :bs + ] + metadata.cache_seqlens_int32.copy_(seq_lens) + num_tokens_per_bs = num_tokens // bs + metadata.cu_seqlens_q = torch.arange( + 0, + bs * num_tokens_per_bs + 1, + num_tokens_per_bs, + dtype=torch.int32, + device=device, + ) + + metadata.cu_seqlens_k = self.draft_extend_metadata["cu_seqlens_k"][ + : (bs + 1) + ] + num_tokens_per_bs = num_tokens // bs + metadata.max_seq_len_q = num_tokens_per_bs + metadata.max_seq_len_k = seq_lens.max().item() + + metadata.page_table = self.draft_extend_metadata["page_table"][:bs, :] + + self.draft_extend_metadata[bs] = metadata + self.forward_metadata = metadata + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + seq_lens_cpu: Optional[torch.Tensor], + ): + """Replay CUDA graph with new inputs.""" + seq_lens = seq_lens[:bs] + seq_lens_cpu = seq_lens_cpu[:bs] + req_pool_indices = req_pool_indices[:bs] + metadata = None + if forward_mode.is_decode_or_idle(): + if spec_info is not None: + # Draft Decode + # Here we only support topk = 1 for now. + metadata = self.decode_cuda_graph_metadata[bs] + max_len = seq_lens_cpu.max().item() + metadata.max_seq_len_k = max_len + self.speculative_step_id + 1 + + max_seq_pages = ( + metadata.max_seq_len_k + self.page_size - 1 + ) // self.page_size + + metadata.cache_seqlens_int32.copy_( + seq_lens + self.speculative_step_id + 1 + ) + else: + # Normal Decode + metadata = self.decode_cuda_graph_metadata[bs] + max_len = seq_lens_cpu.max().item() + max_seq_pages = (max_len + self.page_size - 1) // self.page_size + metadata.max_seq_len_k = max_len + + metadata.cache_seqlens_int32.copy_(seq_lens) + + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32) + ) + page_indices = self.req_to_token[ + req_pool_indices[:, None], + self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages][ + None, : + ], + ] + metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size) + elif forward_mode.is_target_verify(): + # Here we only support topk = 1 for now. + metadata = self.target_verify_metadata[bs] + metadata.cache_seqlens_int32.copy_( + (seq_lens + self.speculative_num_draft_tokens) + ) + + metadata.max_seq_len_k = ( + seq_lens_cpu.max().item() + self.speculative_num_draft_tokens + ) + max_len = seq_lens_cpu.max().item() + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32) + ) + max_seq_pages = ( + metadata.max_seq_len_k + self.page_size - 1 + ) // self.page_size + page_indices = self.req_to_token[ + req_pool_indices[:, None], + self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages], + ] + page_indices //= self.page_size + metadata.page_table[:, :max_seq_pages].copy_(page_indices) + elif forward_mode.is_draft_extend(): + metadata = self.draft_extend_metadata[bs] + metadata.cache_seqlens_int32.copy_(seq_lens) + + metadata.max_seq_len_k = seq_lens_cpu.max().item() + max_len = seq_lens_cpu.max().item() + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32) + ) + accept_length = spec_info.accept_length[:bs] + if spec_info.accept_length_cpu: + metadata.max_seq_len_q = max(spec_info.accept_length_cpu) + 1 + else: + metadata.max_seq_len_q = 1 + + metadata.cu_seqlens_q[1:].copy_( + torch.cumsum(accept_length, dim=0, dtype=torch.int32) + ) + + max_seq_pages = ( + metadata.max_seq_len_k + self.page_size - 1 + ) // self.page_size + page_indices = self.req_to_token[ + req_pool_indices[:, None], + self.draft_extend_metadata["strided_indices"][:max_seq_pages], + ] + metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size) + self.forward_metadata = metadata + + def get_cuda_graph_seq_len_fill_value(self) -> int: + """Get the fill value for sequence lengths in CUDA graph.""" + return 1 + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Initialize the metadata for a forward pass.""" + + metadata = TRTLLMMHAMetadata() + seqlens_in_batch = forward_batch.seq_lens + batch_size = forward_batch.batch_size + device = seqlens_in_batch.device + + if forward_batch.forward_mode.is_decode_or_idle(): + if forward_batch.spec_info is not None: + # Draft Decode + # Here we only support topk = 1 for now. + metadata.cache_seqlens_int32 = ( + seqlens_in_batch + (self.speculative_step_id + 1) + ).to(torch.int32) + metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + ( + self.speculative_step_id + 1 + ) + metadata.cu_seqlens_q = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device=device + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + else: + # Normal Decode + metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32) + metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + metadata.cu_seqlens_q = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device=device + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0) + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + elif forward_batch.forward_mode.is_target_verify(): + # Only support topk = 1 for now. + metadata.cache_seqlens_int32 = ( + forward_batch.seq_lens + self.speculative_num_draft_tokens + ).to(torch.int32) + metadata.max_seq_len_q = self.speculative_num_draft_tokens + metadata.max_seq_len_k = ( + forward_batch.seq_lens_cpu.max().item() + + self.speculative_num_draft_tokens + ) + metadata.cu_seqlens_q = torch.arange( + 0, + batch_size * self.speculative_num_draft_tokens + 1, + self.speculative_num_draft_tokens, + dtype=torch.int32, + device=device, + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32), + (1, 0), + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + + else: + metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32) + metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0) + ) + metadata.page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : metadata.max_seq_len_k + ] + + if ( + any(forward_batch.extend_prefix_lens_cpu) + or forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND + ): + extend_seq_lens = forward_batch.extend_seq_lens + metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu) + metadata.cu_seqlens_q = torch.nn.functional.pad( + torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0) + ) + else: + metadata.max_seq_len_q = metadata.max_seq_len_k + metadata.cu_seqlens_q = metadata.cu_seqlens_k + + # Convert the page table to a strided format + if self.page_size > 1: + self.strided_indices = torch.arange( + 0, metadata.page_table.shape[1], self.page_size, device=self.device + ) + metadata.page_table = ( + metadata.page_table[:, self.strided_indices] // self.page_size + ) + + self.forward_metadata = metadata + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ) -> torch.Tensor: + """Run forward for decode using TRTLLM MHA kernel.""" + cache_loc = forward_batch.out_cache_loc + if save_kv_cache and k is not None: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, layer.k_scale, layer.v_scale + ) + + q = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) + # shape conversion: + # [num_pages, page_size, num_kv_heads, head_dim] -> [num_pages, num_kv_heads, page_size, head_dim] + k_cache = k_cache.view( + -1, self.page_size, layer.tp_k_head_num, layer.head_dim + ).permute(0, 2, 1, 3) + v_cache = v_cache.view( + -1, self.page_size, layer.tp_v_head_num, layer.head_dim + ).permute(0, 2, 1, 3) + kv_cache = (k_cache, v_cache) + + # TODO: add support for quantization + q_scale = 1.0 + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + bmm1_scale = q_scale * k_scale * layer.scaling + bmm2_scale = 1.0 + # sink: additional value per head in the denominator of the softmax. + attention_sink = kwargs.get("sinks", None) + + # Call TRT-LLM kernel + # raw_out: like q, [bs, acc_q_len, num_q_heads, head_dim] but with output dtype + o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=q, + kv_cache=kv_cache, + workspace_buffer=self.workspace_buffer, + block_tables=self.forward_metadata.page_table, + seq_lens=self.forward_metadata.cache_seqlens_int32, + max_seq_len=self.max_context_len, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + window_left=layer.sliding_window_size, + # TODO: add attention_sink operation or nvfp4 scale factor if needed + sinks=attention_sink, + ) + + return o.view(-1, layer.tp_q_head_num * layer.head_dim) + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + **kwargs, + ): + cache_loc = forward_batch.out_cache_loc + if save_kv_cache and k is not None: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, layer.k_scale, layer.v_scale + ) + q = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + # [num_pages, page_size, num_kv_heads, head_dim] -> [num_pages, num_kv_heads, page_size, head_dim] + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) + k_cache = k_cache.view( + -1, self.page_size, layer.tp_k_head_num, layer.head_dim + ).permute(0, 2, 1, 3) + v_cache = v_cache.view( + -1, self.page_size, layer.tp_v_head_num, layer.head_dim + ).permute(0, 2, 1, 3) + kv_cache = (k_cache, v_cache) + + # sink: additional value per head in the denominator of the softmax. + attention_sink = kwargs.get("sinks", None) + # TODO: add support for quantization + q_scale = 1.0 + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + bmm1_scale = q_scale * k_scale * layer.scaling + bmm2_scale = 1.0 + + o = flashinfer.prefill.trtllm_batch_context_with_kv_cache( + query=q, + kv_cache=kv_cache, + workspace_buffer=self.workspace_buffer, + block_tables=self.forward_metadata.page_table, + seq_lens=self.forward_metadata.cache_seqlens_int32, + max_q_len=self.forward_metadata.max_seq_len_q, + max_kv_len=self.max_context_len, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + batch_size=forward_batch.batch_size, + cum_seq_lens_q=self.forward_metadata.cu_seqlens_q, + cum_seq_lens_kv=self.forward_metadata.cu_seqlens_k, + window_left=layer.sliding_window_size, + # TODO: add attention_sink operation or nvfp4 scale factor if needed + sinks=attention_sink, + ) + + return o.view(-1, layer.tp_q_head_num * layer.head_dim) + + +class TRTLLMHAAttnMultiStepDraftBackend(FlashInferMultiStepDraftBackend): + """Multi-step TRTLLM MHA attention kernel used by EAGLE.""" + + def __init__( + self, model_runner: ModelRunner, topk: int, speculative_num_steps: int + ): + super().__init__(model_runner, topk, speculative_num_steps) + for i in range(speculative_num_steps): + self.attn_backends[i] = TRTLLMHAAttnBackend( + model_runner, + skip_prefill=True, + kv_indptr_buf=self.kv_indptr[i], + kv_last_page_len_buf=self.kv_last_page_len, + speculative_step_id=i, + ) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata(forward_batch) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + for i in range(self.speculative_num_steps): + self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens) + + def init_forward_metadata_capture_cuda_graph( + self, + forward_batch: ForwardBatch, + ): + assert forward_batch.spec_info is not None + assert isinstance(forward_batch.spec_info, EagleDraftInput) + + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata_capture_cuda_graph( + forward_batch.batch_size, + forward_batch.batch_size * self.topk, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + encoder_lens=forward_batch.encoder_lens, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + ) + + def init_forward_metadata_replay_cuda_graph( + self, forward_batch: ForwardBatch, bs: int + ): + assert forward_batch.spec_info is not None + assert isinstance(forward_batch.spec_info, EagleDraftInput) + + for i in range(self.speculative_num_steps - 1): + + self.attn_backends[i].init_forward_metadata_replay_cuda_graph( + bs, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + encoder_lens=forward_batch.encoder_lens, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + seq_lens_cpu=forward_batch.seq_lens_cpu, + ) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py new file mode 100755 index 000000000..b8d62c3fa --- /dev/null +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -0,0 +1,589 @@ +from __future__ import annotations + +""" +Support attention backend for TRTLLM MLA kernels from flashinfer. +""" + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Union + +import torch +import triton + +from sglang.srt.layers.attention.flashinfer_mla_backend import ( + FlashInferMLAAttnBackend, + FlashInferMLAMultiStepDraftBackend, +) +from sglang.srt.layers.attention.utils import ( + TRITON_PAD_NUM_PAGE_PER_BLOCK, + create_flashmla_kv_indices_triton, +) +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils import is_flashinfer_available + +if is_flashinfer_available(): + import flashinfer + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpecInfo + +# Constants +DEFAULT_WORKSPACE_SIZE_MB = 128 # Memory workspace size in MB + +# Block constraint from flashinfer requirements +# From flashinfer.decode._check_trtllm_gen_mla_shape: +# block_num % (128 / block_size) == 0 +# This imposes that the total number of blocks must be divisible by +# (128 / block_size). We capture the 128 constant here so we can +# compute the LCM with other padding constraints. +TRTLLM_BLOCK_CONSTRAINT = 128 + +global_zero_init_workspace_buffer = None + + +@dataclass +class TRTLLMMLAPrefillMetadata: + """Metadata for TRTLLM MLA prefill operations.""" + + max_seq_len: int + cum_seq_lens: torch.Tensor + seq_lens: torch.Tensor + + +@dataclass +class TRTLLMMLADecodeMetadata: + """Metadata for TRTLLM MLA decode operations.""" + + block_kv_indices: Optional[torch.Tensor] = None + max_seq_len: Optional[int] = None + + +class TRTLLMMLABackend(FlashInferMLAAttnBackend): + """TRTLLM MLA attention kernel from flashinfer.""" + + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + kv_indptr_buf: Optional[torch.Tensor] = None, + q_indptr_decode_buf: Optional[torch.Tensor] = None, + ): + super().__init__(model_runner, skip_prefill, kv_indptr_buf, q_indptr_decode_buf) + + config = model_runner.model_config + + # Model parameters + self.num_q_heads = config.num_attention_heads // get_attention_tp_size() + self.num_kv_heads = config.get_num_kv_heads(get_attention_tp_size()) + self.num_local_heads = config.num_attention_heads // get_attention_tp_size() + + # MLA-specific dimensions + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim + + # Runtime parameters + self.scaling = config.scaling + self.data_type = model_runner.kv_cache_dtype + self.q_data_type = model_runner.dtype + self.page_size = model_runner.page_size + self.req_to_token = model_runner.req_to_token_pool.req_to_token + + # Workspace allocation + self.workspace_size = DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024 + global global_zero_init_workspace_buffer + if global_zero_init_workspace_buffer is None: + global_zero_init_workspace_buffer = torch.zeros( + self.workspace_size, + dtype=torch.uint8, + device=model_runner.device, + ) + self.workspace_buffer = global_zero_init_workspace_buffer + + # CUDA graph state + self.decode_cuda_graph_metadata = {} + self.decode_cuda_graph_kv_indices = None + self.forward_prefill_metadata: Optional[TRTLLMMLAPrefillMetadata] = None + self.forward_decode_metadata: Union[TRTLLMMLADecodeMetadata, None] = None + + def _calc_padded_blocks(self, max_seq_len: int) -> int: + """ + Calculate padded block count that satisfies both TRT-LLM and Triton constraints. + + Args: + max_seq_len: Maximum sequence length in tokens + + Returns: + Number of blocks padded to satisfy all constraints + """ + blocks = triton.cdiv(max_seq_len, self.page_size) + + # Apply dual constraints (take LCM to satisfy both): + # 1. TRT-LLM: block_num % (128 / page_size) == 0 + # 2. Triton: page table builder uses 64-index bursts, needs multiple of 64 + trtllm_constraint = TRTLLM_BLOCK_CONSTRAINT // self.page_size + constraint_lcm = math.lcm(trtllm_constraint, TRITON_PAD_NUM_PAGE_PER_BLOCK) + + if blocks % constraint_lcm != 0: + blocks = triton.cdiv(blocks, constraint_lcm) * constraint_lcm + return blocks + + def _create_block_kv_indices( + self, + batch_size: int, + max_blocks: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + device: torch.device, + ) -> torch.Tensor: + """ + Create block KV indices tensor using Triton kernel. + + Args: + batch_size: Batch size + max_blocks: Maximum number of blocks per sequence + req_pool_indices: Request pool indices + seq_lens: Sequence lengths + device: Target device + + Returns: + Block KV indices tensor + """ + block_kv_indices = torch.full( + (batch_size, max_blocks), -1, dtype=torch.int32, device=device + ) + + create_flashmla_kv_indices_triton[(batch_size,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + None, + block_kv_indices, + self.req_to_token.stride(0), + max_blocks, + NUM_PAGE_PER_BLOCK=TRITON_PAD_NUM_PAGE_PER_BLOCK, + PAGED_SIZE=self.page_size, + ) + + return block_kv_indices + + def init_cuda_graph_state( + self, + max_bs: int, + max_num_tokens: int, + kv_indices_buf: Optional[torch.Tensor] = None, + ): + """Initialize CUDA graph state for TRTLLM MLA.""" + + max_blocks_per_seq = self._calc_padded_blocks(self.max_context_len) + + self.decode_cuda_graph_kv_indices = torch.full( + (max_bs, max_blocks_per_seq), -1, dtype=torch.int32, device=self.device + ) + + super().init_cuda_graph_state(max_bs, max_num_tokens, kv_indices_buf) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + ): + """Initialize metadata for CUDA graph capture.""" + + # Delegate to parent for non-decode modes. + if not forward_mode.is_decode_or_idle(): + return super().init_forward_metadata_capture_cuda_graph( + bs, + num_tokens, + req_pool_indices, + seq_lens, + encoder_lens, + forward_mode, + spec_info, + ) + + # Custom fast-path for decode/idle. + # Capture with full width so future longer sequences are safe during replay + max_blocks_per_seq = self._calc_padded_blocks(self.max_context_len) + block_kv_indices = self.decode_cuda_graph_kv_indices[:bs, :max_blocks_per_seq] + + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + None, + block_kv_indices, + self.req_to_token.stride(0), + max_blocks_per_seq, + NUM_PAGE_PER_BLOCK=TRITON_PAD_NUM_PAGE_PER_BLOCK, + PAGED_SIZE=self.page_size, + ) + + # Record the true maximum sequence length for this capture batch so that + # the kernel launch path (which requires an int not a tensor) can reuse + # it safely during both capture and replay. + max_seq_len_val = int(seq_lens.max().item()) + + metadata = TRTLLMMLADecodeMetadata( + block_kv_indices, + max_seq_len_val, + ) + self.decode_cuda_graph_metadata[bs] = metadata + self.forward_decode_metadata = metadata + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInfo], + seq_lens_cpu: Optional[torch.Tensor], + ): + """Replay CUDA graph with new inputs.""" + # Delegate to parent for non-decode modes. + if not forward_mode.is_decode_or_idle(): + return super().init_forward_metadata_replay_cuda_graph( + bs, + req_pool_indices, + seq_lens, + seq_lens_sum, + encoder_lens, + forward_mode, + spec_info, + seq_lens_cpu, + ) + + metadata = self.decode_cuda_graph_metadata[bs] + + # Update block indices for new sequences. + create_flashmla_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices[:bs], + seq_lens[:bs], + None, + metadata.block_kv_indices, + self.req_to_token.stride(0), + metadata.block_kv_indices.shape[1], + NUM_PAGE_PER_BLOCK=TRITON_PAD_NUM_PAGE_PER_BLOCK, + PAGED_SIZE=self.page_size, + ) + + # Update stored max_seq_len so subsequent kernel calls use the correct value + # Prefer CPU tensor to avoid GPU synchronization when available. + if seq_lens_cpu is not None: + metadata.max_seq_len = int(seq_lens_cpu.max().item()) + else: + metadata.max_seq_len = int(seq_lens.max().item()) + + def get_cuda_graph_seq_len_fill_value(self) -> int: + """Get the fill value for sequence lengths in CUDA graph.""" + return 1 + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Initialize the metadata for a forward pass.""" + # Delegate to parent for non-decode modes. + if ( + forward_batch.forward_mode.is_extend() + and not forward_batch.forward_mode.is_target_verify() + and not forward_batch.forward_mode.is_draft_extend() + ): + seq_lens = forward_batch.seq_lens - forward_batch.extend_prefix_lens + cum_seq_lens_q = torch.cat( + ( + torch.tensor([0], device=forward_batch.seq_lens.device), + torch.cumsum(seq_lens, dim=0), + ) + ).int() + max_seq_len = max(forward_batch.extend_seq_lens_cpu) + self.forward_prefill_metadata = TRTLLMMLAPrefillMetadata( + max_seq_len, + cum_seq_lens_q, + seq_lens, + ) + elif forward_batch.forward_mode.is_decode_or_idle(): + bs = forward_batch.batch_size + + # Get maximum sequence length. + if getattr(forward_batch, "seq_lens_cpu", None) is not None: + max_seq = forward_batch.seq_lens_cpu.max().item() + else: + max_seq = forward_batch.seq_lens.max().item() + + max_seqlen_pad = self._calc_padded_blocks(max_seq) + block_kv_indices = self._create_block_kv_indices( + bs, + max_seqlen_pad, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens.device, + ) + + max_seq_len_val = int(max_seq) + self.forward_decode_metadata = TRTLLMMLADecodeMetadata( + block_kv_indices, max_seq_len_val + ) + forward_batch.decode_trtllm_mla_metadata = self.forward_decode_metadata + else: + return super().init_forward_metadata(forward_batch) + + def init_mha_chunk_metadata(self, forward_batch: ForwardBatch): + super().init_mha_chunk_metadata(forward_batch, disable_flashinfer_ragged=True) + + def quantize_and_rope_for_fp8( + self, + q_nope: torch.Tensor, + q_rope: torch.Tensor, + k_nope: torch.Tensor, + k_rope: torch.Tensor, + forward_batch: ForwardBatch, + cos_sin_cache: torch.Tensor, + is_neox: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize and apply RoPE for FP8 attention path. + + This function handles the FP8 quantization and RoPE application for MLA attention. + It takes separate query/key nope and rope components, applies RoPE to the rope parts, + quantizes all components to FP8, and merges the query components into a single tensor. + + Args: + q_nope: Query no-position-encoding component [seq_len, num_heads, kv_lora_rank] + - expected dtype: torch.bfloat16 + q_rope: Query RoPE component [seq_len, num_heads, qk_rope_head_dim] + - expected dtype: torch.bfloat16 + k_nope: Key no-position-encoding component [seq_len, num_heads, kv_lora_rank] + - expected dtype: torch.bfloat16 + k_rope: Key RoPE component [seq_len, num_heads, qk_rope_head_dim] + - expected dtype: torch.bfloat16 + forward_batch: Forward batch containing position information + cos_sin_cache: Precomputed cosine/sine cache for RoPE + - expected dtype: matches q_/k_ input dtype (torch.bfloat16) + is_neox: Whether to use NeoX-style RoPE (interleaved) or GPT-style (half rotation) + + Returns: + tuple: (merged_q_out, k_nope_out, k_rope_out) quantized to FP8 + - merged_q_out: [seq_len, num_heads, kv_lora_rank + qk_rope_head_dim], dtype=torch.float8_e4m3fn + - k_nope_out: [seq_len, num_heads, kv_lora_rank], dtype=torch.float8_e4m3fn + - k_rope_out: [seq_len, num_heads, qk_rope_head_dim], dtype=torch.float8_e4m3fn + """ + attn_dtype = torch.float8_e4m3fn + q_len, num_heads = q_rope.shape[0], q_rope.shape[1] + + # Allocate output tensors with FP8 dtype + # Query output will contain merged nope + rope components + q_out = q_rope.new_empty( + q_len, + num_heads, + self.kv_lora_rank + self.qk_rope_head_dim, + dtype=attn_dtype, + ) + + # Key outputs maintain original shapes but with FP8 dtype + k_rope_out = k_rope.new_empty(k_rope.shape, dtype=attn_dtype) + k_nope_out = k_nope.new_empty(k_nope.shape, dtype=attn_dtype) + + # Apply RoPE and quantize all components in a single fused kernel call + # This kernel handles: + # 1. RoPE application to q_rope and k_rope using cos_sin_cache and positions + # 2. Quantization of all components to FP8 format + # 3. Output placement into pre-allocated tensors + flashinfer.rope.mla_rope_quantize_fp8( + q_rope=q_rope, + k_rope=k_rope, + q_nope=q_nope, + k_nope=k_nope, + cos_sin_cache=cos_sin_cache, + pos_ids=forward_batch.positions, + is_neox=is_neox, + quantize_dtype=attn_dtype, + # Output tensor slicing: q_out contains [nope_part, rope_part] + q_rope_out=q_out[..., self.kv_lora_rank :], # RoPE part goes to end + k_rope_out=k_rope_out, + q_nope_out=q_out[..., : self.kv_lora_rank], # Nope part goes to beginning + k_nope_out=k_nope_out, + # Quantization scales (set to 1.0 for no additional scaling) + quant_scale_q=1.0, + quant_scale_kv=1.0, + ) + + return q_out, k_nope_out, k_rope_out + + def forward_decode( + self, + q: torch.Tensor, # q_nope + k: torch.Tensor, # k_nope + v: torch.Tensor, # not used in this backend + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + cos_sin_cache: Optional[torch.Tensor] = None, + is_neox: Optional[bool] = False, + ) -> torch.Tensor: + """Run forward for decode using TRTLLM MLA kernel.""" + merge_query = q_rope is not None + if self.data_type == torch.float8_e4m3fn: + # For FP8 path, we quantize the query and rope parts and merge them into a single tensor + # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend + assert all( + x is not None for x in [q_rope, k_rope, cos_sin_cache] + ), "For FP8 path and using flashinfer.rope.mla_rope_quantize we need all of q_rope, k_rope and cos_sin_cache to be not None." + q, k, k_rope = self.quantize_and_rope_for_fp8( + q, + q_rope, + k.squeeze(1), + k_rope.squeeze(1), + forward_batch, + cos_sin_cache, + is_neox, + ) + merge_query = False + + # Save KV cache if requested + if save_kv_cache: + assert ( + k is not None and k_rope is not None + ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." + forward_batch.token_to_kv_pool.set_mla_kv_buffer( + layer, forward_batch.out_cache_loc, k, k_rope + ) + + # Prepare query tensor inline + if merge_query: + # For FP16 path, we merge the query and rope parts into a single tensor + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope_reshaped = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + query = torch.cat([q_nope, q_rope_reshaped], dim=-1) + else: + # For FP8 path, we already have the query and rope parts merged because of the quantize_and_rope_for_fp8 function + query = q.view(-1, layer.tp_q_head_num, layer.head_dim) + + # Ensure query has shape [bs, acc_q_len, num_q_heads, head_dim] when seq_len 1 + if query.dim() == 3: + query = query.unsqueeze(1) + + # Prepare KV cache inline + k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + kv_cache = k_cache.view(-1, self.page_size, self.kv_cache_dim).unsqueeze(1) + + # Get metadata + metadata = ( + getattr(forward_batch, "decode_trtllm_mla_metadata", None) + or self.forward_decode_metadata + ) + + # Scale computation for TRTLLM MLA kernel BMM1 operation: + # The final BMM1 scale is computed as: q_scale * k_scale * softmax_scale + # Scale components: + # - q_scale: Query scaling factor (set to 1.0 for both FP16/FP8 paths) + # - k_scale: Key scaling factor from model checkpoint (defaults to 1.0 if not available) + # - softmax_scale: Attention softmax scaling = 1/sqrt(head_dim), pre-computed as layer.scaling + # This unified approach works for both FP16 and FP8 quantized attention paths. + q_scale = 1.0 + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + + bmm1_scale = q_scale * k_scale * layer.scaling + + # Call TRT-LLM kernel + raw_out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( + query=query, + kv_cache=kv_cache, + workspace_buffer=self.workspace_buffer, + qk_nope_head_dim=self.qk_nope_head_dim, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + block_tables=metadata.block_kv_indices, + seq_lens=forward_batch.seq_lens.to(torch.int32), + max_seq_len=metadata.max_seq_len, + bmm1_scale=bmm1_scale, + ) + + # Reshape output directly without slicing + output = raw_out.view(-1, layer.tp_q_head_num * layer.v_head_dim) + return output + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend() + ): + return super().forward_extend( + q, k, v, layer, forward_batch, save_kv_cache, q_rope, k_rope + ) + + if not forward_batch.attn_attend_prefix_cache: + q = q.view(-1, layer.tp_q_head_num, layer.head_dim) + k = k.view(-1, layer.tp_k_head_num, layer.head_dim) + v = v.view(-1, layer.tp_k_head_num, layer.v_head_dim) + output = flashinfer.prefill.trtllm_ragged_attention_deepseek( + query=q, + key=k, + value=v, + workspace_buffer=self.workspace_buffer, + seq_lens=self.forward_prefill_metadata.seq_lens, + max_q_len=self.forward_prefill_metadata.max_seq_len, + max_kv_len=self.forward_prefill_metadata.max_seq_len, + bmm1_scale=layer.scaling, + bmm2_scale=1.0, + o_sf_scale=1.0, + batch_size=forward_batch.batch_size, + window_left=-1, + cum_seq_lens_q=self.forward_prefill_metadata.cum_seq_lens, + cum_seq_lens_kv=self.forward_prefill_metadata.cum_seq_lens, + enable_pdl=False, + is_causal=True, + return_lse=forward_batch.mha_return_lse, + ) + else: + # replace with trtllm ragged attention once accuracy is resolved. + output = super().forward_extend( + q, k, v, layer, forward_batch, save_kv_cache, q_rope, k_rope + ) + return output + + +class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend): + """Multi-step draft backend for TRT-LLM MLA used by EAGLE.""" + + def __init__( + self, model_runner: "ModelRunner", topk: int, speculative_num_steps: int + ): + super().__init__(model_runner, topk, speculative_num_steps) + + for i in range(self.speculative_num_steps): + self.attn_backends[i] = TRTLLMMLABackend( + model_runner, + skip_prefill=True, + kv_indptr_buf=self.kv_indptr[i], + q_indptr_decode_buf=self.q_indptr_decode, + ) diff --git a/python/sglang/srt/layers/attention/utils.py b/python/sglang/srt/layers/attention/utils.py new file mode 100644 index 000000000..e8cd2e158 --- /dev/null +++ b/python/sglang/srt/layers/attention/utils.py @@ -0,0 +1,99 @@ +import triton +import triton.language as tl + +# Keep this in sync with the Triton kernel inside `create_flashmla_kv_indices_triton`. +# Number of pages that the kernel writes per iteration. +# Exposed here so other Python modules can import it instead of hard-coding 64. +TRITON_PAD_NUM_PAGE_PER_BLOCK = 64 + + +@triton.jit +def create_flashinfer_kv_indices_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices_ptr, + page_kernel_lens_ptr, + kv_indptr, + kv_start_idx, + kv_indices_ptr, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(axis=0) + + # find the req pool idx, this is for batch to token + req_pool_index = tl.load(req_pool_indices_ptr + pid) + kv_indices_offset = tl.load(kv_indptr + pid) + + kv_start = 0 + kv_end = 0 + if kv_start_idx: + kv_start = tl.load(kv_start_idx + pid).to(tl.int32) + kv_end = kv_start + kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) + + num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + for i in range(num_loop): + # index into req_to_token_ptr needs to be int64 + offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE + mask = offset < kv_end - kv_start + data = tl.load( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + kv_start + + offset, + mask=mask, + ) + tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask) + + +@triton.jit +def create_flashmla_kv_indices_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices_ptr, + page_kernel_lens_ptr, + kv_start_idx, + kv_indices_ptr, + req_to_token_ptr_stride: tl.constexpr, + kv_indices_ptr_stride: tl.constexpr, + NUM_PAGE_PER_BLOCK: tl.constexpr = TRITON_PAD_NUM_PAGE_PER_BLOCK, + PAGED_SIZE: tl.constexpr = 64, +): + BLOCK_SIZE: tl.constexpr = 4096 + pid = tl.program_id(axis=0) + + # find the req pool idx, this is for batch to token + req_pool_index = tl.load(req_pool_indices_ptr + pid) + + kv_start = 0 + kv_end = 0 + if kv_start_idx: + kv_start = tl.load(kv_start_idx + pid).to(tl.int32) + kv_end = kv_start + + kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) + + num_paged = tl.cdiv(kv_end - kv_start, PAGED_SIZE) + num_pages_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + + for i in range(num_pages_loop): + # index into req_to_token_ptr needs to be int64 + paged_offset = ( + tl.arange(0, NUM_PAGE_PER_BLOCK).to(tl.int64) + i * NUM_PAGE_PER_BLOCK + ) * PAGED_SIZE + paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK + + mask = paged_offset < num_paged * PAGED_SIZE + mask_out = paged_offset_out < num_paged + + data = tl.load( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + kv_start + + paged_offset, + mask=mask, + ) + tl.store( + kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out, + data // PAGED_SIZE, + mask=mask_out, + ) diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py new file mode 100644 index 000000000..2be3e450b --- /dev/null +++ b/python/sglang/srt/layers/attention/vision.py @@ -0,0 +1,630 @@ +from __future__ import annotations + +import dataclasses +import functools +import math +from functools import lru_cache, partial +from typing import Any, Callable, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size +from sglang.srt.utils import ( + get_device_capability, + is_blackwell, + is_cuda, + print_info_once, +) + +_is_cuda = is_cuda() + +if _is_cuda: + from sgl_kernel.flash_attn import flash_attn_varlen_func + +from sglang.srt.distributed import ( + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, +) +from sglang.srt.distributed import utils as dist_utils +from sglang.srt.layers.attention.triton_ops.prefill_attention import ( + context_attention_fwd, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.utils import add_prefix + +ROTARY_EMBED_CLASSES = { + "normal": apply_rotary_pos_emb, +} + + +@dataclasses.dataclass +class SingletonCache: + data: Any = None + + def set_data(self, value: Any) -> None: + self.data = value + + def get_data(self) -> Optional[Any]: + return self.data + + def empty(self) -> bool: + return self.get_data() is None + + +# TODO: requires real seqlens from images +@functools.lru_cache(maxsize=128) +def _get_cu_seqlens_for_shape(batch_size: int, seqlen: int, device) -> torch.Tensor: + """ + Generates cumulative sequence lengths (cu_seqlens) for a given batch_size, seqlen, and device. + Caches the result based on these parameters. + """ + cu_seqlens = torch.arange( + 0, + (batch_size + 1) * seqlen, + step=seqlen, + dtype=torch.int32, + device=device, + ) + return cu_seqlens + + +class VisionSdpaAttention(nn.Module): + r""" + Scaled Dot Product Attention inner product + + """ + + def __init__( + self, + head_dim: int, + num_heads: int, + num_kv_heads: int, + dropout: float = 0.0, + flatten_batch: bool = False, + softmax_in_single_precision: bool = False, + **kwargs, + ): + super().__init__() + self.head_size = head_dim + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.flatten_batch = flatten_batch + self.softmax_in_single_precision = softmax_in_single_precision + self.dropout = dropout + self.scale = 1.0 / math.sqrt(self.head_size) + + @staticmethod + @lru_cache(maxsize=128) + def _generate_mask_cache( + s: int, flatten_batch: bool, cu_seqlens: tuple + ) -> torch.BoolTensor: + """ + Generate a boolean attention mask with caching mechanism. + Args: + s: sequence length + flatten_batch: whether to flatten batch dimension + cu_seqlens: tuple of cumulative sequence lengths + Returns: + attention mask tensor of shape [b, 1, s, s] or [1, s, s] + """ + if flatten_batch: + mask = torch.zeros([1, s, s], dtype=torch.bool) + for i in range(1, len(cu_seqlens)): + start = cu_seqlens[i - 1] + end = cu_seqlens[i] + mask[..., start:end, start:end] = True + else: + # [1, 1, 1, s] + row_indices = torch.arange(s).view(1, 1, 1, s) + # [1, 1, s, 1] + col_indices = torch.arange(s).view(1, 1, s, 1) + # [b, 1, 1, 1] + seq_lens = torch.tensor( + [end - start for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:])], + ).view(-1, 1, 1, 1) + + mask = (row_indices < seq_lens) & (col_indices < seq_lens) + + return mask + + def generate_patch_attention_mask( + self, + s: int, + cu_seqlens: Optional[torch.Tensor], + flatten_batch: bool = False, + ) -> Optional[torch.Tensor]: + r""" + Creates a non-causal 4D mask of shape `(b, 1, s, s)` or `(1, 1, s, s)`. + Args: + s: sequence length + cu_seqlens: cumulative sequence lengths tensor. If not, returns an empty mask + flatten_batch: whether to flatten batch dimension + Returns: + attention mask tensor or None + """ + if cu_seqlens is None: + return None + + cu_seqlens_tuple = tuple(cu_seqlens.cpu().tolist()) + + return self._generate_mask_cache(s, flatten_batch, cu_seqlens_tuple) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + bsz: int, + cu_seqlens: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + r""" + Args: + cu_seqlens: [b] + Returns: + [b * s, h, head_size] + """ + if self.flatten_batch: + assert bsz == 1, "flatten_batch is True, bsz must be 1" + + assert q.dim() == 3, q.shape + + s = q.shape[0] // bsz + + # [b, 1, s, s] + if attention_mask is None: + attention_mask = self.generate_patch_attention_mask( + s, cu_seqlens, flatten_batch=self.flatten_batch + ) + + if attention_mask is None: + if self.softmax_in_single_precision: + raise RuntimeError("Empty attention mask") + else: + attention_mask = attention_mask.to(device=q.device) + + q, k, v = [rearrange(x, "(b s) h d -> b h s d", b=bsz) for x in [q, k, v]] + + if self.softmax_in_single_precision: + k = rearrange(k, "b h s d -> b h d s") + attn_weights = torch.matmul(q, k) * self.scale + del k + # masking + attention_mask = (~attention_mask) * torch.finfo(q.dtype).min + attn_weights = attn_weights + attention_mask + del attention_mask + # full-precision + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(q.dtype) + attn_weights = nn.functional.dropout( + attn_weights, p=self.dropout, training=False + ) + output = torch.matmul(attn_weights, v) + del attn_weights, v + else: + # SDPA + # [b, h, s, head_size] + output = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attention_mask, + dropout_p=self.dropout, + is_causal=False, + ) + + # [b, h, s, head_size] --> [b * s, h, head_size] + output = rearrange(output, "b h s d -> (b s) h d") + + return output + + +class VisionTritonAttention(nn.Module): + """ + Triton-implemented attention without a causal mask + """ + + def __init__( + self, + **kwargs, + ): + super().__init__() + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + bsz: int, + seq_len: int, + **kwargs, + ) -> torch.Tensor: + r""" + Args: + cu_seqlens: [b] + Returns: + [b * s, h, head_size] + """ + if cu_seqlens is None: + cu_seqlens = _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device) + + # [b * s, head, head_size] + output = torch.empty_like(q) + seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] + max_seqlen = seq_lens.max().item() + context_attention_fwd( + q, + k, + v, + output, + cu_seqlens.cuda(), + seq_lens.cuda(), + max_seqlen, + is_causal=False, + ) + + return output + + +class VisionFlash3Attention(nn.Module): + def __init__( + self, + **kwargs, + ): + if not _is_cuda: + raise Exception("VisionFlash3Attention is only available for cuda") + super().__init__() + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens: Optional[Union[SingletonCache, torch.Tensor]], + bsz: int, + seq_len: int, + **kwargs, + ) -> torch.Tensor: + r""" + Args: + cu_seqlens: [b] + Returns: + [b * s, h, head_size] + """ + if cu_seqlens is None: + cu_seqlens = _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device) + elif isinstance(cu_seqlens, SingletonCache): + if cu_seqlens.empty(): + cu_seqlens.set_data( + _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device) + ) + cu_seqlens = cu_seqlens.get_data() + + cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device) + seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] + max_seqlen = seq_lens.max().item() + + output = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + ) + + return output + + +QKV_BACKEND_IMPL = { + "triton_attn": VisionTritonAttention, + "sdpa": VisionSdpaAttention, + "fa3": VisionFlash3Attention, +} + + +class VisionAttention(nn.Module): + r""" + Multi-headed attention without any cache, mostly used for multimodal transformers. + + + Args: + use_qkv_parallel (bool, optional): If True, use QKV-parallel attention. + softmax_in_single_precision (bool, default to False): + if ``True``, the softmax will be performed in single-precision + Otherwise, it will be performed in half-precision + + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + projection_size: int, + use_qkv_parallel: bool, + qkv_backend: Optional[str] = None, + quant_config: Optional[QuantizationConfig] = None, + dropout: float = 0.0, + softmax_in_single_precision: bool = False, + flatten_batch: bool = False, + prefix: str = "", + proj_bias: bool = True, + num_dummy_heads: int = 0, + qkv_bias: bool = True, + qk_normalization: bool = False, + layer_norm_eps: float = 1e-06, + customized_position_embedding_applier: Callable[ + [torch.Tensor, torch.Tensor, Any, Any], Tuple[torch.Tensor, torch.Tensor] + ] = None, + **kwargs, + ): + super().__init__() + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + self.tp_size = attn_tp_size + self.tp_rank = attn_tp_rank + self.dropout = dropout + self.head_size = embed_dim // num_heads + self.hidden_size_per_attention_head = dist_utils.divide( + projection_size, num_heads + ) + self.num_attention_heads_per_partition = dist_utils.divide( + num_dummy_heads + num_heads, self.tp_size + ) + self.num_attention_kv_heads_per_partition = dist_utils.divide( + num_dummy_heads + num_heads, self.tp_size + ) + + self.q_size = self.num_attention_heads_per_partition * self.head_size + self.kv_size = self.num_attention_kv_heads_per_partition * self.head_size + + self.qk_normalization = qk_normalization + + # Additional dummy heads are used to enable TP for common GPU counts. + self.dummy_dim = (num_dummy_heads + num_heads) * self.head_size + + if self.qk_normalization: + self.q_norm = RMSNorm( + self.dummy_dim, eps=layer_norm_eps, var_hidden_size=embed_dim + ) + self.k_norm = RMSNorm( + self.dummy_dim, eps=layer_norm_eps, var_hidden_size=embed_dim + ) + + # Select attention backend via a unified method + _passed_backend = qkv_backend + qkv_backend = self._determine_attention_backend(_passed_backend) + if ( + global_server_args_dict["mm_attention_backend"] is None + and _passed_backend is None + ): + print_info_once(f"Multimodal attention backend not set. Use {qkv_backend}.") + print_info_once(f"Using {qkv_backend} as multimodal attention backend.") + + self.customized_position_embedding_applier = ( + customized_position_embedding_applier + ) + self.qkv_backend = QKV_BACKEND_IMPL[qkv_backend]( + head_dim=self.head_size, + num_heads=self.num_attention_heads_per_partition, + num_kv_heads=self.num_attention_kv_heads_per_partition, + dropout=dropout, + flatten_batch=flatten_batch, + softmax_in_single_precision=softmax_in_single_precision, + ) + + self.use_qkv_parallel = use_qkv_parallel + if use_qkv_parallel: + self.qkv_proj = QKVParallelLinear( + hidden_size=embed_dim, + head_size=self.head_size, + total_num_heads=num_dummy_heads + num_heads, + total_num_kv_heads=num_dummy_heads + num_heads, + bias=qkv_bias, + quant_config=quant_config, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + prefix=add_prefix("qkv_proj", prefix), + ) + else: + self.qkv_proj = ColumnParallelLinear( + input_size=embed_dim, + output_size=3 * self.dummy_dim, + bias=qkv_bias, + quant_config=quant_config, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + prefix=add_prefix("qkv_proj", prefix), + ) + self.proj = RowParallelLinear( + input_size=self.dummy_dim, + output_size=embed_dim, + bias=proj_bias, + quant_config=quant_config, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + prefix=add_prefix("proj", prefix), + ) + + def _determine_attention_backend(self, passed_backend: Optional[str]) -> str: + """Decide the multimodal attention backend string. + + Priority: server args override > constructor arg > platform default. + + Platform defaults: + - CUDA: "triton_attn" + - Non-CUDA: "sdpa" + """ + override_backend = global_server_args_dict["mm_attention_backend"] + if override_backend is not None: + backend = override_backend + elif passed_backend is not None: + backend = passed_backend + elif is_cuda(): + major, minor = get_device_capability() + if major == 9: + backend = "fa3" + else: + backend = "triton_attn" + else: + backend = "sdpa" + if backend == "fa3" and is_blackwell(): + raise ValueError("The 'fa3' backend is not supported on Blackwell GPUs") + + return backend + + def _apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor): + """apply qk norm for internvl vit attn""" + q = q.flatten(1, 2) + k = k.flatten(1, 2) + + if self.tp_size > 1: + q = tensor_model_parallel_all_gather(q.contiguous()) + k = tensor_model_parallel_all_gather(k.contiguous()) + q = self.q_norm(q) + k = self.k_norm(k) + if self.tp_size > 1: + splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size) + q = splitter(q)[self.tp_rank] + k = splitter(k)[self.tp_rank] + q = q.unflatten(-1, (-1, self.head_size)) + k = k.unflatten(-1, (-1, self.head_size)) + return q, k + + def forward( + self, + x: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + r""" + Args: + x: [b, s, embed_dim] + cu_seqlens: [b] + Returns: + [s, b, head * head_size] + """ + if x.dim() == 2: + x = x.unsqueeze(0) + assert x.dim() == 3, x.shape + x_shape = x.shape + bsz, s, _ = x_shape + head = self.num_attention_heads_per_partition + kv_head = self.num_attention_kv_heads_per_partition + if self.use_qkv_parallel: + # [b, s, embed_dim] --> [b, s, embed_dim] + qkv, _ = self.qkv_proj(x) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + # [b, s, embed_dim] --> [b * s, head, head_size] + q = q.reshape(bsz * s, head, -1).contiguous() + k = k.reshape(bsz * s, kv_head, -1).contiguous() + v = v.reshape(bsz * s, kv_head, -1).contiguous() + else: + # [b, s, embed_dim] --> [s, b, embed_dim] + x = rearrange(x, "b s ... -> s b ...") + # [s, b, embed_dim] --> [s, b, head * 3 * head_size] + qkv, _ = self.qkv_proj(x) + + # [s, b, head, head_dim_sum] + new_x_shape = qkv.size()[:-1] + ( + head, + self.q_size + 2 * self.kv_size, + ) + qkv = qkv.view(*new_x_shape) + + # [s, b, head, 3 * head_size] --> 3 [s, b, head, head_size] + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + # [s, b, head, head_size] --> [b, s, head, head_size] + q, k, v = [ + rearrange(x, "s b ... -> b s ...").contiguous() for x in (q, k, v) + ] + + if position_embeddings is not None: + original_shape = q.shape + + if self.customized_position_embedding_applier is not None: + q, k = self.customized_position_embedding_applier( + q, k, position_embeddings, x_shape + ) + q = q.view(original_shape) + k = k.view(original_shape) + else: + cos, sin = position_embeddings + + # [total_tokens, head, head_size] + q = q.view(-1, head, self.head_size) + k = k.view(-1, head, self.head_size) + + q, k = apply_rotary_pos_emb(q, k, cos, sin) + + q = q.view(original_shape) + k = k.view(original_shape) + + if q.dim() == 4: + # [b, s, head, head_size] --> [b * s, head, head_size] + q = rearrange(q, "b s ... -> (b s) ...") + if k.dim() == 4: + # [b, s, head, head_size] --> [b * s, head, head_size] + k = rearrange(k, "b s ... -> (b s) ...") + if v.dim() == 4: + # [b, s, head, head_size] --> [b * s, head, head_size] + v = rearrange(v, "b s ... -> (b s) ...") + + assert q.dim() == 3, q.dim() + assert k.dim() == 3, k.dim() + assert v.dim() == 3, v.dim() + + # internvl + if self.qk_normalization: + q, k = self._apply_qk_norm(q, k) + + output = self.qkv_backend.forward( + q=q, + k=k, + v=v, + bsz=bsz, + seq_len=s, + cu_seqlens=cu_seqlens, + attention_mask=attention_mask, + ) + + assert output.dim() == 3, output.shape + + if self.use_qkv_parallel: + # [b * s, h, head_size] --> [b, s, h * head_size] + output = rearrange(output, "(b s) ... h d -> b s ... (h d)", b=bsz) + + # [b, s, h * head_size] --> [b, s, h * head_size] + output, _ = self.proj(output) + else: + # [b * s, h, head_size] --> [s, b, h * head_size] + context_layer = rearrange( + output, "(b s) h d -> s b (h d)", b=bsz, s=s + ).contiguous() + + # [s, b, h * head_size] --> [s, b, h * head_size] + output, _ = self.proj(context_layer) + + # [s, b, h * head_size] --> [b, s, h * head_size] + output = output.view(bsz, s, -1) + + return output diff --git a/python/sglang/srt/layers/attention/vision_utils.py b/python/sglang/srt/layers/attention/vision_utils.py new file mode 100644 index 000000000..ecccb1f85 --- /dev/null +++ b/python/sglang/srt/layers/attention/vision_utils.py @@ -0,0 +1,65 @@ +"""Utility functions for vision attention layers.""" + +import torch + +from sglang.srt.layers.dp_attention import get_attention_tp_size + + +def update_vit_attn_dummy_heads_config(config): + """Update HF config to ensure vision attention num_attention_heads is divisible by tp_size""" + tp_size = get_attention_tp_size() + num_heads = getattr( + config.vision_config, + "num_heads", + getattr(config.vision_config, "num_attention_heads", None), + ) + head_dim = config.vision_config.hidden_size // num_heads + num_dummy_heads = 0 + + if num_heads % tp_size != 0: + num_dummy_heads = ((num_heads + tp_size - 1) // tp_size) * tp_size - num_heads + + setattr(config.vision_config, "head_dim", head_dim) + setattr(config.vision_config, "num_dummy_heads", num_dummy_heads) + + +def pad_vit_attn_dummy_heads(config, name: str, loaded_weight: torch.Tensor): + """Pad attention qkv weights for dummy heads""" + num_dummy_heads = config.vision_config.num_dummy_heads + if num_dummy_heads == 0: + return loaded_weight + head_dim = config.vision_config.head_dim + + if "attn.qkv_proj" in name: + wq, wk, wv = loaded_weight.chunk(3, dim=0) + if name.endswith(".weight"): + dummy_shape = [num_dummy_heads, head_dim, wq.shape[-1]] + elif name.endswith(".bias"): + dummy_shape = [num_dummy_heads, head_dim] + else: + raise RuntimeError(f"Unsupported weight with name={name}") + pad_func = lambda x: torch.cat( + [x.unflatten(0, (-1, head_dim)), x.new_zeros(dummy_shape)], dim=0 + ).flatten(0, 1) + wq, wk, wv = pad_func(wq), pad_func(wk), pad_func(wv) + loaded_weight = torch.cat([wq, wk, wv], dim=0) + elif any([_ in name for _ in ["attn.q_proj", "attn.k_proj", "attn.v_proj"]]): + if name.endswith(".weight"): + dummy_shape = [num_dummy_heads, head_dim, loaded_weight.shape[-1]] + elif name.endswith(".bias"): + dummy_shape = [num_dummy_heads, head_dim] + else: + raise RuntimeError(f"Unsupported weight with name={name}") + padded_weight = loaded_weight.new_zeros(dummy_shape) + loaded_weight = torch.cat( + [loaded_weight.unflatten(0, (-1, head_dim)), padded_weight], dim=0 + ).flatten(0, 1) + elif "attn.proj.weight" in name: + padded_weight = loaded_weight.new_zeros( + loaded_weight.shape[0], head_dim * num_dummy_heads + ) + loaded_weight = torch.cat([loaded_weight, padded_weight], dim=-1) + elif "attn.q_norm.weight" in name or "attn.k_norm.weight" in name: + padded_weight = loaded_weight.new_zeros(head_dim * num_dummy_heads) + loaded_weight = torch.cat([loaded_weight, padded_weight], dim=0) + return loaded_weight diff --git a/python/sglang/srt/layers/attention/wave_backend.py b/python/sglang/srt/layers/attention/wave_backend.py new file mode 100644 index 000000000..eb6e061ac --- /dev/null +++ b/python/sglang/srt/layers/attention/wave_backend.py @@ -0,0 +1,627 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Union + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils import get_bool_env_var, get_device_core_count + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput + +logger = logging.getLogger(__name__) + + +@triton.jit +def get_num_kv_splits_triton( + num_kv_splits_ptr, + seq_lens_ptr, + num_seq, + num_group, + num_head, + num_kv_head, + max_kv_splits, + device_core_count, + MAX_NUM_SEQ: tl.constexpr, +): + # TODO: this method is tunable, we need more online serving data to tune it + offs_seq = tl.arange(0, MAX_NUM_SEQ) + mask_seq = offs_seq < num_seq + + seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=0) + max_seq_len = tl.max(seq_lens) + seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=max_seq_len) + min_seq_len = tl.min(seq_lens) + if max_seq_len * 8 < min_seq_len * 10: + min_seq_len = max_seq_len + max_kv_splits_1 = tl.minimum(tl.cdiv(max_seq_len, min_seq_len), max_kv_splits) + kv_chunk_size_1 = tl.cdiv(max_seq_len, max_kv_splits_1) + + # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually + ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0 + ext_device_core_count = tl.cast( + device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32 + ) + block_h, num_kv_group = 16, num_head // num_kv_head + if num_kv_group == 1: + token_grid = num_seq * num_group * num_head + else: + # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd + block_h = tl.minimum(block_h, num_kv_group) + token_grid = num_seq * num_group * tl.cdiv(num_head, block_h) + max_kv_splits_2 = tl.minimum( + tl.cdiv(ext_device_core_count, token_grid), max_kv_splits + ) + kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2) + + num_kv_splits = tl.maximum( + tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2) + ) + + offs_token = offs_seq * num_group + mask_token = offs_token < num_seq * num_group + for i in range(0, num_group): + tl.store(num_kv_splits_ptr + i + offs_token, num_kv_splits, mask=mask_token) + + +@dataclass +class ForwardMetadata: + attn_logits: torch.Tensor + attn_lse: torch.Tensor + max_extend_len: int + num_kv_splits: torch.Tensor + kv_indptr: torch.Tensor + kv_indices: torch.Tensor + qo_indptr: torch.Tensor + custom_mask: torch.Tensor + mask_indptr: torch.Tensor + + +class WaveAttnBackend(AttentionBackend): + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + kv_indptr_buf: Optional[torch.Tensor] = None, + ): + # Lazy import to avoid the initialization of cuda context + from sglang.srt.layers.attention.wave_ops.decode_attention import ( + decode_attention_fwd, + ) + from sglang.srt.layers.attention.wave_ops.extend_attention import ( + extend_attention_wave, + ) + + super().__init__() + + # Set unique cache dir for each process to avoid cache write races + import wave_lang.kernel.wave.cache as cache + + base_cache_dir = cache.CACHE_BASE_DIR + new_dir = base_cache_dir / f"worker_{model_runner.tp_rank}" + logger.info(f"Setting Wave cache dir: {new_dir}") + cache.CACHE_BASE_DIR = new_dir + + self.decode_attention_fwd = decode_attention_fwd + self.extend_attention_fwd = extend_attention_wave + + self.skip_prefill = skip_prefill + + max_bs = model_runner.req_to_token_pool.size + + if kv_indptr_buf is None: + self.kv_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + else: + self.kv_indptr = kv_indptr_buf + + self.req_to_token = model_runner.req_to_token_pool.req_to_token + + if not self.skip_prefill: + self.qo_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + + self.mask_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int64, device=model_runner.device + ) + + self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens + + self.num_head = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.num_kv_head = model_runner.model_config.get_num_kv_heads( + get_attention_tp_size() + ) + + self.static_kv_splits = get_bool_env_var( + "SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS", "false" + ) + self.max_kv_splits = model_runner.server_args.triton_attention_num_kv_splits + self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1] + + self.forward_metadata: ForwardMetadata = None + + self.max_context_len = model_runner.model_config.context_len + + self.device = model_runner.device + self.device_core_count = get_device_core_count(model_runner.gpu_id) + + def get_num_kv_splits( + self, + num_kv_splits: torch.Tensor, + seq_lens: torch.Tensor, + ): + num_token, num_seq = num_kv_splits.shape[0], seq_lens.shape[0] + num_group = num_token // num_seq + + assert ( + num_group * num_seq == num_token + ), f"num_seq({num_seq}), num_token({num_token}), something goes wrong!" + + if self.static_kv_splits or self.device_core_count <= 0: + num_kv_splits.fill_(self.max_kv_splits) + return + + if num_seq < 256: + SCHEDULE_SEQ = 256 + else: + SCHEDULE_SEQ = triton.next_power_of_2(num_seq) + + get_num_kv_splits_triton[(1,)]( + num_kv_splits, + seq_lens, + num_seq, + num_group, + self.num_head, + self.num_kv_head, + self.max_kv_splits, + self.device_core_count, + MAX_NUM_SEQ=SCHEDULE_SEQ, + ) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init auxiliary variables for wave attention backend.""" + + bs = forward_batch.batch_size + kv_indptr = self.kv_indptr + spec_info = forward_batch.spec_info + + if forward_batch.forward_mode.is_decode_or_idle(): + if spec_info is None: + kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + forward_batch.seq_lens_sum, dtype=torch.int32, device=self.device + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + else: + kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices + bs = kv_indptr.shape[0] - 1 + + from sglang.srt.layers.attention.wave_ops.decode_attention import ( + decode_attention_intermediate_arrays_shapes, + ) + + attn_logits_shape, attn_logits_max_shape = ( + decode_attention_intermediate_arrays_shapes( + bs, self.v_head_dim, self.num_head, self.max_kv_splits + ) + ) + attn_logits = torch.empty( + attn_logits_shape, + dtype=torch.float32, + device=self.device, + ) + attn_lse = torch.empty( + attn_logits_max_shape, + dtype=torch.float32, + device=self.device, + ) + num_kv_splits = torch.empty((bs,), dtype=torch.int32, device=self.device) + + self.get_num_kv_splits(num_kv_splits, forward_batch.seq_lens) + + qo_indptr = None + custom_mask = None + mask_indptr = None + max_extend_len = None + elif forward_batch.forward_mode.is_target_verify(): + bs = len(forward_batch.req_pool_indices) + qo_indptr = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + # Different with flashinfer kv_indptr and kv_indices construction + kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + kv_indptr[-1], dtype=torch.int32, device=self.device + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + + custom_mask = spec_info.custom_mask + seq_mask_len = self.num_draft_tokens * ( + forward_batch.seq_lens + self.num_draft_tokens + ) + mask_indptr = self.mask_indptr + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0) + mask_indptr = mask_indptr[: bs + 1] + max_extend_len = self.num_draft_tokens + num_kv_splits = None + attn_logits = None + attn_lse = None + elif forward_batch.forward_mode.is_draft_extend(): + kv_indices, kv_indptr, qo_indptr, custom_mask = ( + spec_info.generate_attn_arg_prefill( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + None, + self.req_to_token, + ) + ) + mask_indptr = None + # TODO(FIXME): This will trigger an invalid Eagle tree when using + # `max(spec_info.accept_length_cpu)`. + # It might have been forgotten to update somewhere. + max_extend_len = torch.max(spec_info.accept_length).item() + num_kv_splits = None + attn_logits = None + attn_lse = None + else: + kv_indptr[1 : bs + 1] = torch.cumsum( + forward_batch.extend_prefix_lens, dim=0 + ) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + forward_batch.extend_prefix_lens.sum().item(), + dtype=torch.int32, + device=self.device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.extend_prefix_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + + qo_indptr = self.qo_indptr + qo_indptr[1 : bs + 1] = torch.cumsum(forward_batch.extend_seq_lens, dim=0) + qo_indptr = qo_indptr[: bs + 1] + custom_mask = None + mask_indptr = None + attn_logits = None + attn_lse = None + max_extend_len = torch.max(forward_batch.extend_seq_lens).item() + num_kv_splits = None + + self.forward_metadata = ForwardMetadata( + attn_logits, + attn_lse, + max_extend_len, + num_kv_splits, + kv_indptr, + kv_indices, + qo_indptr, + custom_mask, + mask_indptr, + ) + + def init_cuda_graph_state( + self, + max_bs: int, + max_num_tokens: int, + kv_indices_buf: Optional[torch.Tensor] = None, + ): + from sglang.srt.layers.attention.wave_ops.decode_attention import ( + decode_attention_intermediate_arrays_shapes, + ) + + attn_logits_shape, attn_logits_max_shape = ( + decode_attention_intermediate_arrays_shapes( + max_bs, self.v_head_dim, self.num_head, self.max_kv_splits + ) + ) + self.cuda_graph_attn_logits = torch.zeros( + attn_logits_shape, + dtype=torch.float32, + device=self.device, + ) + self.cuda_graph_attn_lse = torch.zeros( + attn_logits_max_shape, + dtype=torch.float32, + device=self.device, + ) + self.cuda_graph_num_kv_splits = torch.full( + (max_bs,), self.max_kv_splits, dtype=torch.int32, device=self.device + ) + if kv_indices_buf is None: + self.cuda_graph_kv_indices = torch.zeros( + (max_bs * self.max_context_len), + dtype=torch.int32, + device=self.device, + ) + else: + self.cuda_graph_kv_indices = kv_indices_buf + + if not self.skip_prefill: + self.cuda_graph_custom_mask = torch.zeros( + (max_bs * self.max_context_len), + dtype=torch.uint8, + device=self.device, + ) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + assert encoder_lens is None, "Not supported" + + if forward_mode.is_decode_or_idle(): + if spec_info is None: + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + else: + kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices + + attn_logits = self.cuda_graph_attn_logits + attn_lse = self.cuda_graph_attn_lse + max_extend_len = None + num_kv_splits = self.cuda_graph_num_kv_splits + qo_indptr = None + custom_mask = None + mask_indptr = None + elif forward_mode.is_target_verify(): + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + + custom_mask = self.cuda_graph_custom_mask + seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens) + mask_indptr = self.mask_indptr[: bs + 1] + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0) + max_extend_len = self.num_draft_tokens + num_kv_splits = None + attn_logits = None + attn_lse = None + else: + raise ValueError( + f"Invalid forward mode: {forward_mode=} for CUDA Graph capture." + ) + + self.forward_metadata = ForwardMetadata( + attn_logits, + attn_lse, + max_extend_len, + num_kv_splits, + kv_indptr, + kv_indices, + qo_indptr, + custom_mask, + mask_indptr, + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + seq_lens_cpu: Optional[torch.Tensor], + ): + # NOTE: encoder_lens expected to be zeros or None + if forward_mode.is_decode_or_idle(): + # Update kv_indptr, kv_indices + kv_indptr = self.kv_indptr + kv_indices = self.cuda_graph_kv_indices + num_kv_splits = self.cuda_graph_num_kv_splits + if spec_info is None: + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens[:bs], dim=0) + kv_indptr = kv_indptr[: bs + 1] + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices[:bs], + seq_lens[:bs], + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + num_token = bs + else: + kv_indptr[: spec_info.kv_indptr.shape[0]] = spec_info.kv_indptr + kv_indices[: spec_info.kv_indices.shape[0]] = spec_info.kv_indices + num_token = spec_info.kv_indptr.shape[0] - 1 + self.get_num_kv_splits(num_kv_splits[:num_token], seq_lens[:bs]) + elif forward_mode.is_target_verify(): + # Update qo_indptr, kv_indptr, kv_indices, custom_mask, mask_indptr + bs = len(req_pool_indices) + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + custom_mask = self.cuda_graph_custom_mask + custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask + seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens) + mask_indptr = self.mask_indptr[: bs + 1] + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0) + else: + raise ValueError( + f"Invalid forward mode: {forward_mode=} for CUDA Graph replay." + ) + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + # TODO: reuse the buffer across layers + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + + max_extend_len = self.forward_metadata.max_extend_len + computed_max_ext_seq_len = torch.max(forward_batch.extend_seq_lens) + if computed_max_ext_seq_len != max_extend_len: + assert len(forward_batch.extend_seq_lens) == 1 + forward_batch.extend_seq_lens[0] = max_extend_len + forward_batch.seq_lens = max_extend_len + + self.extend_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + k.contiguous(), + v.contiguous(), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + self.forward_metadata.qo_indptr, + self.forward_metadata.kv_indptr, + self.forward_metadata.kv_indices, + self.forward_metadata.custom_mask, + self.forward_metadata.mask_indptr, + self.forward_metadata.max_extend_len, + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + is_causal=True, + layer_scaling=layer.scaling, + logit_cap=layer.logit_cap, + ) + return o + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + ): + # During torch.compile, there is a bug in rotary_emb that causes the + # output value to have a 3D tensor shape. This reshapes the output correctly. + q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim) + + # TODO: reuse the buffer across layers + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + + self.decode_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + self.forward_metadata.kv_indptr, + self.forward_metadata.kv_indices, + self.forward_metadata.attn_logits, + self.forward_metadata.attn_lse, + self.forward_metadata.num_kv_splits, + self.max_kv_splits, + layer.scaling, + layer.logit_cap, + ) + return o diff --git a/python/sglang/srt/layers/attention/wave_ops/decode_attention.py b/python/sglang/srt/layers/attention/wave_ops/decode_attention.py new file mode 100644 index 000000000..c76bee9af --- /dev/null +++ b/python/sglang/srt/layers/attention/wave_ops/decode_attention.py @@ -0,0 +1,184 @@ +""" +Memory-efficient attention for decoding. +It supports page size = 1. +""" + +import functools +import logging + +from wave_lang.kernel.lang.global_symbols import * +from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile +from wave_lang.kernel.wave.constraints import GenericDot, MMAOperand, MMAType +from wave_lang.kernel.wave.templates.paged_decode_attention import ( + get_paged_decode_attention_kernels, + get_paged_decode_intermediate_arrays_shapes, + paged_decode_attention_shape, +) +from wave_lang.kernel.wave.utils.general_utils import get_default_scheduling_params +from wave_lang.kernel.wave.utils.run_utils import set_default_run_config + +logger = logging.getLogger(__name__) +import os + +dump_generated_mlir = int(os.environ.get("WAVE_DUMP_MLIR", 0)) + + +@functools.lru_cache(maxsize=4096) +def get_wave_kernel( + shape: paged_decode_attention_shape, + max_kv_splits, + input_dtype, + output_dtype, + logit_cap, +): + mha = (shape.num_query_heads // shape.num_kv_heads) == 1 + + # Get the kernels (either compile or load from cache). + if mha: + mfma_variant = ( + GenericDot(along_dim=MMAOperand.M, k_vec_size=4, k_mult=1), + GenericDot(along_dim=MMAOperand.M, k_vec_size=1, k_mult=64), + ) + else: + mfma_variant = (MMAType.F32_16x16x16_F16, MMAType.F32_16x16x16_F16) + + ( + phase_0, + phase_1, + hyperparams_0, + hyperparams_1, + dynamic_symbols_0, + dynamic_symbols_1, + ) = get_paged_decode_attention_kernels( + shape, + mfma_variant, + max_kv_splits, + input_dtype=input_dtype, + output_dtype=output_dtype, + logit_cap=logit_cap, + ) + hyperparams_0.update(get_default_scheduling_params()) + hyperparams_1.update(get_default_scheduling_params()) + + options = WaveCompileOptions( + subs=hyperparams_0, + canonicalize=True, + run_bench=False, + use_buffer_ops=True, + waves_per_eu=2, + dynamic_symbols=dynamic_symbols_0, + wave_runtime=True, + ) + options = set_default_run_config(options) + phase_0 = wave_compile(options, phase_0) + + options = WaveCompileOptions( + subs=hyperparams_1, + canonicalize=True, + run_bench=False, + use_buffer_ops=False, + waves_per_eu=4, + dynamic_symbols=dynamic_symbols_1, + wave_runtime=True, + ) + options = set_default_run_config(options) + phase_1 = wave_compile(options, phase_1) + + return phase_0, phase_1 + + +def decode_attention_intermediate_arrays_shapes( + num_seqs, head_size_kv, num_query_heads, max_kv_splits +): + # Not all fields are used, but we need to pass them to the function + shape = paged_decode_attention_shape( + num_query_heads=num_query_heads, + num_kv_heads=0, + head_size=0, + head_size_kv=head_size_kv, + block_size=0, + num_seqs=num_seqs, + ) + return get_paged_decode_intermediate_arrays_shapes(shape, max_kv_splits) + + +def decode_attention_wave( + q, + k_buffer, + v_buffer, + o, + b_req_idx, + req_to_token, + attn_logits, + attn_logits_max, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap, +): + num_seqs, num_query_heads, head_size = q.shape + _, num_kv_heads, _ = k_buffer.shape + _, _, head_size_kv = v_buffer.shape + block_size = 32 + shape = paged_decode_attention_shape( + num_query_heads, + num_kv_heads, + head_size, + head_size_kv, + block_size, + num_seqs, + ) + + phase_0, phase_1 = get_wave_kernel( + shape, max_kv_splits, q.dtype, o.dtype, logit_cap + ) + + mb_qk = phase_0( + q, + k_buffer, + v_buffer, + b_req_idx, + req_to_token, + attn_logits, + attn_logits_max, + ) + if dump_generated_mlir: + filename = f"wave_decode_attention_phase0_{'x'.join(map(str, shape))}.mlir" + with open(filename, "w") as f: + f.write(mb_qk.module_op.get_asm()) + + mb_sv = phase_1(attn_logits, attn_logits_max, b_req_idx, o) + if dump_generated_mlir: + filename = f"wave_decode_attention_phase1_{'x'.join(map(str, shape))}.mlir" + with open(filename, "w") as f: + f.write(mb_sv.module_op.get_asm()) + + +def decode_attention_fwd( + q, + k_buffer, + v_buffer, + o, + b_req_idx, + req_to_token, + attn_logits, + attn_logits_max, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap=0.0, +): + decode_attention_wave( + q, + k_buffer, + v_buffer, + o, + b_req_idx, + req_to_token, + attn_logits, + attn_logits_max, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap, + ) diff --git a/python/sglang/srt/layers/attention/wave_ops/extend_attention.py b/python/sglang/srt/layers/attention/wave_ops/extend_attention.py new file mode 100644 index 000000000..27e674db2 --- /dev/null +++ b/python/sglang/srt/layers/attention/wave_ops/extend_attention.py @@ -0,0 +1,147 @@ +""" +Memory-efficient attention for prefill. +It support page size = 1. +""" + +import functools +import os + +import torch +from wave_lang.kernel.lang.global_symbols import * +from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile +from wave_lang.kernel.wave.constraints import MMAType +from wave_lang.kernel.wave.scheduling.schedule import SchedulingType +from wave_lang.kernel.wave.templates.attention_common import AttentionShape +from wave_lang.kernel.wave.templates.extend_attention import get_extend_attention_kernel +from wave_lang.kernel.wave.utils.general_utils import get_default_scheduling_params +from wave_lang.kernel.wave.utils.run_utils import set_default_run_config + +dump_generated_mlir = int(os.environ.get("WAVE_DUMP_MLIR", 0)) + + +@functools.lru_cache +def get_wave_kernel( + shape: AttentionShape, + q_shape: tuple[int], + k_shape: tuple[int], + v_shape: tuple[int], + k_cache_shape: tuple[int], + v_cache_shape: tuple[int], + o_shape: tuple[int], + input_dtype: torch.dtype, + output_dtype: torch.dtype, + size_dtype: torch.dtype, + is_causal: bool, + logit_cap: float, + layer_scaling: float, +): + assert shape.num_query_heads % shape.num_kv_heads == 0 + + mfma_variant = (MMAType.F32_16x16x32_K8_F16, MMAType.F32_16x16x16_F16) + ( + extend_attention, + hyperparams, + dynamic_symbols, + ) = get_extend_attention_kernel( + shape, + mfma_variant, + q_shape, + k_shape, + v_shape, + k_cache_shape, + v_cache_shape, + o_shape, + input_dtype=input_dtype, + output_dtype=output_dtype, + size_dtype=size_dtype, + is_causal=is_causal, + layer_scaling=layer_scaling, + logit_cap=logit_cap, + ) + + hyperparams.update(get_default_scheduling_params()) + options = WaveCompileOptions( + subs=hyperparams, + canonicalize=True, + run_bench=False, + schedule=SchedulingType.NONE, + use_scheduling_barriers=False, + dynamic_symbols=dynamic_symbols, + use_buffer_ops=True, + waves_per_eu=2, + denorm_fp_math_f32="preserve-sign", + wave_runtime=True, + ) + options = set_default_run_config(options) + extend_attention = wave_compile(options, extend_attention) + + return extend_attention + + +def extend_attention_wave( + q_extend, + k_extend, + v_extend, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + custom_mask, + mask_indptr, + max_seq_len, + output, + is_causal=True, + layer_scaling=None, + logit_cap=0, +): + shape = AttentionShape( + num_query_heads=q_extend.shape[1], + num_kv_heads=k_extend.shape[1], + head_size=q_extend.shape[2], + head_size_kv=k_extend.shape[2], + num_seqs=kv_indptr.shape[0] - 1, + max_seq_len=max_seq_len, + ) + + # Run the wave kernel. + extend_attention = get_wave_kernel( + shape, + q_extend.shape, + k_extend.shape, + v_extend.shape, + k_buffer.shape, + v_buffer.shape, + output.shape, + input_dtype=q_extend.dtype, + output_dtype=output.dtype, + size_dtype=qo_indptr.dtype, + is_causal=is_causal, + layer_scaling=layer_scaling, + logit_cap=logit_cap, + ) + + mb = extend_attention( + q_extend, + k_extend, + v_extend, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + max_seq_len, + output, + ) + + if dump_generated_mlir: + shape_list = [ + q_extend.shape[0], + q_extend.shape[1], + k_extend.shape[1], + q_extend.shape[2], + k_extend.shape[2], + ] + filename = f"wave_prefill_attention_{'x'.join(map(str, shape_list))}.mlir" + with open(filename, "w") as f: + f.write(mb.module_op.get_asm()) diff --git a/python/sglang/srt/layers/attention/wave_ops/prefill_attention.py b/python/sglang/srt/layers/attention/wave_ops/prefill_attention.py new file mode 100644 index 000000000..2d8aa4678 --- /dev/null +++ b/python/sglang/srt/layers/attention/wave_ops/prefill_attention.py @@ -0,0 +1,79 @@ +""" +Memory-efficient attention for prefill. +It support page size = 1. +""" + +import math +import os + +from wave_lang.kernel.lang.global_symbols import * +from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile +from wave_lang.kernel.wave.constraints import MMAType +from wave_lang.kernel.wave.templates.attention_common import AttentionShape +from wave_lang.kernel.wave.templates.prefill_attention import ( + get_prefill_attention_kernel, +) +from wave_lang.kernel.wave.utils.general_utils import get_default_scheduling_params +from wave_lang.kernel.wave.utils.run_utils import set_default_run_config + +dump_generated_mlir = int(os.environ.get("WAVE_DUMP_MLIR", 0)) + + +def prefill_attention_wave( + q, k, v, o, b_start_loc, b_seq_len, max_seq_len, is_causal=True +): + + shape = AttentionShape( + num_query_heads=q.shape[1], + num_kv_heads=k.shape[1], + head_size=q.shape[2], + head_size_kv=k.shape[2], + num_seqs=b_seq_len.shape[0], + max_seq_len=max_seq_len, + total_seq_len=q.shape[0], + ) + + assert shape.num_query_heads % shape.num_kv_heads == 0 + + output_shape = (shape.total_seq_len, shape.num_query_heads, shape.head_size_kv) + # Run the wave kernel. + mfma_variant = (MMAType.F32_16x16x16_F16, MMAType.F32_16x16x16_F16) + (prefill, hyperparams) = get_prefill_attention_kernel( + shape, + mfma_variant, + q.shape, + k.shape, + v.shape, + output_shape, + input_dtype=q.dtype, + output_dtype=o.dtype, + size_dtype=b_seq_len.dtype, + ) + + hyperparams.update(get_default_scheduling_params()) + + log2e = 1.44269504089 + dk_sqrt = math.sqrt(1.0 / shape.head_size) + + options = WaveCompileOptions( + subs=hyperparams, + canonicalize=True, + run_bench=False, + use_scheduling_barriers=False, + ) + options = set_default_run_config(options) + prefill = wave_compile(options, prefill) + + mb = prefill( + q * dk_sqrt * log2e, + k, + v, + b_start_loc, + b_seq_len, + o, + ) + if dump_generated_mlir: + shape_list = [q.shape[0], q.shape[1], k.shape[1], q.shape[2], k.shape[2]] + filename = f"wave_prefill_attention_{'x'.join(map(str, shape_list))}.mlir" + with open(filename, "w") as f: + f.write(mb.module_op.get_asm()) diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py new file mode 100644 index 000000000..fba8d8f18 --- /dev/null +++ b/python/sglang/srt/layers/communicator.py @@ -0,0 +1,676 @@ +# Copyright 2023-2024 SGLang 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 dataclasses import dataclass +from enum import Enum, auto +from functools import partial +from typing import Dict, Optional + +import torch + +from sglang.srt.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.dp_attention import ( + attn_tp_all_gather_into_tensor, + attn_tp_reduce_scatter_tensor, + dp_gather_partial, + dp_reduce_scatter_tensor, + dp_scatter, + get_attention_dp_size, + get_attention_tp_rank, + get_attention_tp_size, + get_global_dp_buffer, + get_local_dp_buffer, + is_dp_attention_enabled, +) +from sglang.srt.layers.moe import ( + get_moe_a2a_backend, + should_use_flashinfer_cutlass_moe_fp4_allgather, +) +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.utils import ( + get_bool_env_var, + is_cuda, + is_flashinfer_available, + is_gfx95_supported, + is_hip, + is_sm90_supported, + is_sm100_supported, +) + +_is_flashinfer_available = is_flashinfer_available() +_is_sm90_supported = is_cuda() and is_sm90_supported() +_is_sm100_supported = is_cuda() and is_sm100_supported() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() +_is_gfx95_supported = is_gfx95_supported() + +if _use_aiter and _is_gfx95_supported: + from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant + +FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048 + + +class ScatterMode(Enum): + """ + Suppose we have TP=4, DP=2, enable-dp-attention, and the system handles seq a,b,c,d + Model input/output: [ab, ab, cd, cd] for four ranks respectively + SCATTERED: [a, b, c, d] + TP_ATTN_FULL: [ab, ab, cd, cd], i.e. all ranks inside a TP attn group have full data of the group + FULL: [abcd, abcd, abcd, abcd] + """ + + SCATTERED = auto() + TP_ATTN_FULL = auto() + FULL = auto() + + @staticmethod + def model_input_output(): + """The scatter mode for model forward pass input and output data""" + return ScatterMode.TP_ATTN_FULL + + +@dataclass +class _LayerModeComputationContext: + num_layers: int + layer_id: int + is_layer_sparse: bool + is_previous_layer_sparse: Optional[bool] + + def previous_layer(self): + assert self.is_previous_layer_sparse is not None + return _LayerModeComputationContext( + layer_id=self.layer_id - 1, + is_layer_sparse=self.is_previous_layer_sparse, + is_previous_layer_sparse=None, + num_layers=self.num_layers, + ) + + +@dataclass +class LayerScatterModes: + layer_input_mode: ScatterMode + attn_mode: ScatterMode + # Can be further split into e.g. mlp_input_mode and mlp_output_mode if needed + mlp_mode: ScatterMode + middle_residual_mode: ScatterMode + layer_output_mode: ScatterMode + + @classmethod + def init_new(cls, **kwargs): + context = _LayerModeComputationContext(**kwargs) + return cls( + layer_input_mode=cls._compute_layer_input_mode(context), + attn_mode=ScatterMode.TP_ATTN_FULL, + mlp_mode=cls._compute_mlp_mode(context), + middle_residual_mode=cls._compute_middle_residual_mode(context), + layer_output_mode=cls._compute_layer_output_mode(context), + ) + + @classmethod + def _compute_layer_input_mode(cls, context: _LayerModeComputationContext): + if context.layer_id == 0: + return ScatterMode.model_input_output() + return cls._compute_layer_output_mode(context.previous_layer()) + + @classmethod + def _compute_mlp_mode(cls, context: _LayerModeComputationContext): + if context.is_layer_sparse: + return ( + ScatterMode.SCATTERED + if ( + # Token dispatch/combine will be handled outside of LayerCommunicator for these modes. + not get_moe_a2a_backend().is_none() + or should_use_flashinfer_cutlass_moe_fp4_allgather() + ) + else ScatterMode.FULL + ) + else: + return ( + ScatterMode.SCATTERED + if enable_moe_dense_fully_dp() + else ScatterMode.FULL + ) + + @classmethod + def _compute_middle_residual_mode(cls, context: _LayerModeComputationContext): + mlp_mode = cls._compute_mlp_mode(context) + if mlp_mode == ScatterMode.SCATTERED: + return ScatterMode.SCATTERED + if mlp_mode == ScatterMode.FULL: + return ScatterMode.TP_ATTN_FULL + raise NotImplementedError + + @classmethod + def _compute_layer_output_mode(cls, context: _LayerModeComputationContext): + mlp_mode = cls._compute_mlp_mode(context) + if context.layer_id == context.num_layers - 1: + return ScatterMode.model_input_output() + if mlp_mode == ScatterMode.SCATTERED: + return ScatterMode.SCATTERED + if mlp_mode == ScatterMode.FULL: + return ScatterMode.TP_ATTN_FULL + raise NotImplementedError + + +def enable_moe_dense_fully_dp(): + return global_server_args_dict["moe_dense_tp_size"] == 1 + + +class LayerCommunicator: + def __init__( + self, + layer_scatter_modes: LayerScatterModes, + input_layernorm: torch.nn.Module, + post_attention_layernorm: torch.nn.Module, + # Reduce scatter requires skipping all-reduce in model code after MoE/MLP, so only enable for models which have that implemented. Remove flag once done for all models that use LayerCommunicator. + allow_reduce_scatter: bool = False, + is_last_layer: bool = False, + ): + self.layer_scatter_modes = layer_scatter_modes + self.input_layernorm = input_layernorm + self.post_attention_layernorm = post_attention_layernorm + self.allow_reduce_scatter = allow_reduce_scatter + self.is_last_layer = is_last_layer + + self._context = CommunicateContext.init_new() + self._communicate_simple_fn = CommunicateSimpleFn.get_fn( + input_mode=self.layer_scatter_modes.layer_input_mode, + output_mode=self.layer_scatter_modes.attn_mode, + context=self._context, + ) + self._communicate_with_all_reduce_and_layer_norm_fn = ( + CommunicateWithAllReduceAndLayerNormFn.get_fn( + hidden_states_input_mode=self.layer_scatter_modes.attn_mode, + residual_input_mode=self.layer_scatter_modes.layer_input_mode, + hidden_states_output_mode=self.layer_scatter_modes.mlp_mode, + residual_output_mode=self.layer_scatter_modes.middle_residual_mode, + context=self._context, + ) + ) + self._communicate_summable_tensor_pair_fn = ( + CommunicateSummableTensorPairFn.get_fn( + hidden_states_input_mode=self.layer_scatter_modes.mlp_mode, + residual_input_mode=self.layer_scatter_modes.middle_residual_mode, + output_mode=self.layer_scatter_modes.layer_output_mode, + context=self._context, + ) + ) + + def prepare_attn( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + qaunt_format: str = "", + ): + if hidden_states.shape[0] == 0: + residual = hidden_states + else: + if ( + residual is not None + and hasattr(hidden_states, "_sglang_needs_allreduce_fusion") + and hidden_states._sglang_needs_allreduce_fusion + ): + hidden_states, residual = ( + self.input_layernorm.forward_with_allreduce_fusion( + hidden_states, residual + ) + ) + else: + if residual is None: + residual = hidden_states + + if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format): + hidden_states = fused_rms_mxfp4_quant( + hidden_states, + self.input_layernorm.weight, + self.input_layernorm.variance_epsilon, + None, + None, + None, + None, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + else: + if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format): + hidden_states, residual = fused_rms_mxfp4_quant( + hidden_states, + self.input_layernorm.weight, + self.input_layernorm.variance_epsilon, + None, + None, + None, + residual, + ) + else: + hidden_states, residual = self.input_layernorm( + hidden_states, residual + ) + + hidden_states = self._communicate_simple_fn( + hidden_states=hidden_states, + forward_batch=forward_batch, + context=self._context, + ) + + return hidden_states, residual + + def prepare_mlp( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + ): + return self._communicate_with_all_reduce_and_layer_norm_fn( + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + layernorm=self.post_attention_layernorm, + context=self._context, + ) + + def postprocess_layer( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + ): + return self._communicate_summable_tensor_pair_fn( + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + context=self._context, + allow_reduce_scatter=self.allow_reduce_scatter, + ) + + def should_use_reduce_scatter(self, forward_batch: ForwardBatch): + return ( + self.allow_reduce_scatter + and self._communicate_summable_tensor_pair_fn + is CommunicateSummableTensorPairFn._scatter_hidden_states + and forward_batch.dp_padding_mode.is_max_len() + ) + + def should_fuse_mlp_allreduce_with_next_layer( + self, forward_batch: ForwardBatch + ) -> bool: + speculative_algo = global_server_args_dict.get("speculative_algorithm", None) + if ( + is_dp_attention_enabled() + and speculative_algo is not None + and speculative_algo.is_eagle() + ): + return False + + batch_size = ( + forward_batch.input_ids.shape[0] + if hasattr(forward_batch, "input_ids") + else 0 + ) + if batch_size > FUSE_ALLREDUCE_MAX_BATCH_SIZE: + return False + + static_conditions_met = ( + (not self.is_last_layer) + and (self._context.tp_size > 1) + and global_server_args_dict.get("enable_flashinfer_allreduce_fusion", False) + and _is_flashinfer_available + ) + + if not static_conditions_met: + return False + + return ( + batch_size > 0 + and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE + and (not self.is_last_layer) + ) + + +@dataclass +class CommunicateContext: + process_group_sizes: Dict[ScatterMode, int] + attn_tp_rank: int + attn_tp_size: int + attn_dp_size: int + tp_size: int + + def is_same_group_size(self, a: ScatterMode, b: ScatterMode): + return self.process_group_sizes[a] == self.process_group_sizes[b] + + @classmethod + def init_new(cls): + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + attn_dp_size = get_attention_dp_size() + tp_size = get_tensor_model_parallel_world_size() + process_group_sizes = { + ScatterMode.SCATTERED: 1, + ScatterMode.TP_ATTN_FULL: attn_tp_size, + # TODO: support --moe-dense-tp-size > 1 + ScatterMode.FULL: tp_size, + } + return cls( + process_group_sizes=process_group_sizes, + attn_tp_rank=attn_tp_rank, + attn_tp_size=attn_tp_size, + attn_dp_size=attn_dp_size, + tp_size=tp_size, + ) + + +class CommunicateSimpleFn: + @staticmethod + def get_fn( + input_mode: ScatterMode, + output_mode: ScatterMode, + context: CommunicateContext, + ): + if context.is_same_group_size(input_mode, output_mode): + return CommunicateSimpleFn._trivial + + if (input_mode == ScatterMode.SCATTERED) and ( + output_mode == ScatterMode.TP_ATTN_FULL + ): + return CommunicateSimpleFn._scattered_to_tp_attn_full + + raise NotImplementedError(f"{input_mode=} {output_mode=}") + + @staticmethod + def _trivial( + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + ) -> torch.Tensor: + return hidden_states + + @staticmethod + def _scattered_to_tp_attn_full( + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + ) -> torch.Tensor: + hidden_states, local_hidden_states = ( + get_local_dp_buffer(), + hidden_states, + ) + attn_tp_all_gather_into_tensor( + hidden_states, + local_hidden_states, + ) + return hidden_states + + +class CommunicateWithAllReduceAndLayerNormFn: + """Besides communication, needs to + 1. All reduce in tp_attn_group on hidden_states + 2. Apply layer norm + """ + + @staticmethod + def get_fn( + hidden_states_input_mode: ScatterMode, + residual_input_mode: ScatterMode, + hidden_states_output_mode: ScatterMode, + residual_output_mode: ScatterMode, + context: CommunicateContext, + ): + + if ( + context.is_same_group_size( + hidden_states_input_mode, hidden_states_output_mode + ) + and context.is_same_group_size(residual_input_mode, residual_output_mode) + and context.attn_tp_size == 1 + ): + return CommunicateWithAllReduceAndLayerNormFn._simple + + if ( + (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL) + and ( + residual_input_mode in [ScatterMode.SCATTERED, ScatterMode.TP_ATTN_FULL] + ) + and (hidden_states_output_mode == ScatterMode.FULL) + and (residual_output_mode == ScatterMode.TP_ATTN_FULL) + ): + return partial( + CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual, + residual_input_mode=residual_input_mode, + ) + + if ( + (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL) + and ( + residual_input_mode in [ScatterMode.SCATTERED, ScatterMode.TP_ATTN_FULL] + ) + and (hidden_states_output_mode == ScatterMode.SCATTERED) + and (residual_output_mode == ScatterMode.SCATTERED) + ): + return partial( + CommunicateWithAllReduceAndLayerNormFn._scatter_hidden_states_and_residual, + residual_input_mode=residual_input_mode, + ) + + raise NotImplementedError( + f"{hidden_states_input_mode=} {residual_input_mode=} {hidden_states_output_mode=} {residual_output_mode=}" + ) + + @staticmethod + def _simple( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: CommunicateContext, + ): + # TODO move these `if shape != 0` into LayerNorm itself + if hidden_states.shape[0] != 0: + hidden_states, residual = layernorm(hidden_states, residual) + return hidden_states, residual + + @staticmethod + def _gather_hidden_states_and_residual( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: CommunicateContext, + *, + residual_input_mode, + ): + if residual_input_mode == ScatterMode.SCATTERED and context.attn_tp_size > 1: + residual, local_residual = ( + get_local_dp_buffer(), + residual, + ) + attn_tp_all_gather_into_tensor(residual, local_residual) + if context.attn_dp_size != 1: + if context.attn_tp_rank == 0: + hidden_states += residual + + # Perform layernorm on smaller data before comm. Only valid when attn_tp_size is 1 (tp_size == dp_size) + use_layer_norm_before_gather = context.attn_tp_size == 1 + if use_layer_norm_before_gather and hidden_states.shape[0] != 0: + residual = hidden_states + hidden_states = layernorm(hidden_states) + hidden_states, local_hidden_states = ( + get_global_dp_buffer(), + hidden_states, + ) + dp_gather_partial(hidden_states, local_hidden_states, forward_batch) + + if not use_layer_norm_before_gather: + dp_scatter(residual, hidden_states, forward_batch) + if hidden_states.shape[0] != 0: + hidden_states = layernorm(hidden_states) + else: + # According to the discussion in https://github.com/flashinfer-ai/flashinfer/issues/1223#issuecomment-3047256465 + # We set the max token num to 128 for allreduce fusion with min-latency case(use_oneshot=True). + if ( + (_is_sm100_supported or _is_sm90_supported) + and _is_flashinfer_available + and hasattr(layernorm, "forward_with_allreduce_fusion") + and global_server_args_dict["enable_flashinfer_allreduce_fusion"] + and hidden_states.shape[0] <= 4096 + ): + hidden_states, residual = layernorm.forward_with_allreduce_fusion( + hidden_states, residual + ) + else: + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + hidden_states, residual = layernorm(hidden_states, residual) + return hidden_states, residual + + @staticmethod + def _scatter_hidden_states_and_residual( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: CommunicateContext, + *, + residual_input_mode, + ): + input_hidden_states = hidden_states + hidden_states = hidden_states.tensor_split(context.attn_tp_size)[ + context.attn_tp_rank + ] + attn_tp_reduce_scatter_tensor(hidden_states, input_hidden_states) + if residual_input_mode == ScatterMode.TP_ATTN_FULL: + residual = residual.tensor_split(context.attn_tp_size)[context.attn_tp_rank] + if hidden_states.shape[0] != 0: + hidden_states, residual = layernorm(hidden_states, residual) + return hidden_states, residual + + +class CommunicateSummableTensorPairFn: + """It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed.""" + + @classmethod + def execute( + cls, + hidden_states_input_mode, + residual_input_mode, + output_mode, + context, + **kwargs, + ): + return cls.get_fn( + hidden_states_input_mode=hidden_states_input_mode, + residual_input_mode=residual_input_mode, + output_mode=output_mode, + context=context, + )(context=context, **kwargs) + + @staticmethod + def get_fn( + hidden_states_input_mode: ScatterMode, + residual_input_mode: ScatterMode, + output_mode: ScatterMode, + context: CommunicateContext, + ): + if context.is_same_group_size( + hidden_states_input_mode, output_mode + ) and context.is_same_group_size(residual_input_mode, output_mode): + return CommunicateSummableTensorPairFn._trivial + + if ( + (hidden_states_input_mode == ScatterMode.FULL) + and (residual_input_mode == ScatterMode.TP_ATTN_FULL) + and (output_mode == ScatterMode.TP_ATTN_FULL) + ): + return CommunicateSummableTensorPairFn._scatter_hidden_states + + if ( + (hidden_states_input_mode == ScatterMode.SCATTERED) + and (residual_input_mode == ScatterMode.SCATTERED) + and (output_mode == ScatterMode.TP_ATTN_FULL) + ): + return CommunicateSummableTensorPairFn._gather + + if ( + (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL) + and (residual_input_mode == ScatterMode.TP_ATTN_FULL) + and (output_mode == ScatterMode.SCATTERED) + ): + return CommunicateSummableTensorPairFn._scatter + + raise NotImplementedError( + f"{hidden_states_input_mode=} {residual_input_mode=} {output_mode=}" + ) + + @staticmethod + def _trivial( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + **kwargs, + ): + return hidden_states, residual + + @staticmethod + def _scatter_hidden_states( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + allow_reduce_scatter: bool = False, + ): + hidden_states, global_hidden_states = ( + get_local_dp_buffer(), + hidden_states, + ) + if allow_reduce_scatter and forward_batch.dp_padding_mode.is_max_len(): + # When using padding, all_reduce is skipped after MLP and MOE and reduce scatter is used here instead. + dp_reduce_scatter_tensor(hidden_states, global_hidden_states) + else: + dp_scatter(hidden_states, global_hidden_states, forward_batch) + return hidden_states, residual + + @staticmethod + def _gather( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + **kwargs, + ): + hidden_states += residual + residual = None + hidden_states, local_hidden_states = ( + get_local_dp_buffer(), + hidden_states, + ) + attn_tp_all_gather_into_tensor( + hidden_states, + local_hidden_states, + ) + return hidden_states, residual + + @staticmethod + def _scatter( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + ): + assert residual is None, "not yet handled residual!=None" + tensor_list = list(hidden_states.tensor_split(context.attn_tp_size)) + hidden_states = tensor_list[context.attn_tp_rank] + return hidden_states, residual diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py new file mode 100644 index 000000000..1250636eb --- /dev/null +++ b/python/sglang/srt/layers/dp_attention.py @@ -0,0 +1,489 @@ +from __future__ import annotations + +import functools +import logging +from contextlib import contextmanager +from enum import IntEnum, auto +from typing import TYPE_CHECKING, List, Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.distributed import ( + GroupCoordinator, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, + tensor_model_parallel_all_reduce, +) + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +_ATTN_TP_GROUP: Optional[GroupCoordinator] = None +_ATTN_TP_RANK: Optional[int] = None +_ATTN_TP_SIZE: Optional[int] = None +_ATTN_DP_RANK: Optional[int] = None +_ATTN_DP_SIZE: Optional[int] = None +_LOCAL_ATTN_DP_SIZE: Optional[int] = None +_LOCAL_ATTN_DP_RANK: Optional[int] = None +_ENABLE_DP_ATTENTION_FLAG: bool = False + + +class DpPaddingMode(IntEnum): + + # Padding tokens to max length and then gather tokens using `all_gather_into_tensor` + MAX_LEN = auto() + # Padding tokens to sum length and then gather tokens using `all_reduce` + SUM_LEN = auto() + + def is_max_len(self): + return self == DpPaddingMode.MAX_LEN + + def is_sum_len(self): + return self == DpPaddingMode.SUM_LEN + + @classmethod + def get_dp_padding_mode(cls, global_num_tokens: List[int]) -> DpPaddingMode: + # we choose the mode that minimizes the communication cost + max_len = max(global_num_tokens) + sum_len = sum(global_num_tokens) + if sum_len * 2 > max_len * get_attention_dp_size(): + return cls.MAX_LEN + else: + return cls.SUM_LEN + + @classmethod + def get_default_mode_in_cuda_graph(cls) -> DpPaddingMode: + return cls.MAX_LEN + + +class _DpGatheredBufferWrapper: + + _hidden_size: int + _dtype: torch.dtype + _device: torch.device + _global_dp_buffer_len: int + _local_dp_buffer_len: int + _global_num_tokens: Optional[List[int]] + + @classmethod + def set_metadata(cls, hidden_size: int, dtype: torch.dtype, device: torch.device): + cls._hidden_size = hidden_size + cls._dtype = dtype + cls._device = device + + @classmethod + def set_dp_buffer_len( + cls, + global_dp_buffer_len: int, + local_dp_buffer_len: int, + global_num_tokens: Optional[List[int]] = None, + ): + cls._global_dp_buffer_len = global_dp_buffer_len + cls._local_dp_buffer_len = local_dp_buffer_len + cls._global_num_tokens = global_num_tokens + + @classmethod + def get_global_dp_buffer(cls) -> torch.Tensor: + return torch.empty( + (cls._global_dp_buffer_len, cls._hidden_size), + dtype=cls._dtype, + device=cls._device, + ) + + @classmethod + def get_local_dp_buffer(cls) -> torch.Tensor: + return torch.empty( + (cls._local_dp_buffer_len, cls._hidden_size), + dtype=cls._dtype, + device=cls._device, + ) + + @classmethod + def get_global_dp_buffer_len(cls) -> int: + return cls._global_dp_buffer_len + + @classmethod + def get_local_dp_buffer_len(cls) -> int: + return cls._local_dp_buffer_len + + @classmethod + def get_dp_global_num_tokens(cls) -> List[int]: + return cls._global_num_tokens + + +def set_dp_buffer_len( + global_dp_buffer_len: int, + local_dp_buffer_len: int, + global_num_tokens: Optional[List[int]] = None, +): + _DpGatheredBufferWrapper.set_dp_buffer_len( + global_dp_buffer_len, local_dp_buffer_len, global_num_tokens + ) + + +def get_global_dp_buffer() -> torch.Tensor: + return _DpGatheredBufferWrapper.get_global_dp_buffer() + + +def get_local_dp_buffer() -> torch.Tensor: + return _DpGatheredBufferWrapper.get_local_dp_buffer() + + +def get_global_dp_buffer_len() -> int: + return _DpGatheredBufferWrapper.get_global_dp_buffer_len() + + +def get_local_dp_buffer_len() -> int: + return _DpGatheredBufferWrapper.get_local_dp_buffer_len() + + +def get_dp_global_num_tokens() -> List[int]: + return _DpGatheredBufferWrapper.get_dp_global_num_tokens() + + +def compute_dp_attention_world_info(enable_dp_attention, tp_rank, tp_size, dp_size): + if not enable_dp_attention: + return tp_rank, tp_size, 0 + + attn_tp_size = tp_size // dp_size + attn_dp_rank = tp_rank // attn_tp_size + attn_tp_rank = tp_rank % attn_tp_size + + return attn_tp_rank, attn_tp_size, attn_dp_rank + + +def compute_dp_attention_local_info( + enable_dp_attention, tp_rank, tp_size, dp_size, moe_dense_tp_size +): + if not enable_dp_attention: + return tp_rank, tp_size, 0 + + local_tp_size = moe_dense_tp_size if moe_dense_tp_size else tp_size + local_tp_rank = tp_rank % local_tp_size + local_dp_size = max(1, dp_size // (tp_size // local_tp_size)) + + local_attn_tp_size = local_tp_size // local_dp_size + local_attn_dp_rank = local_tp_rank // local_attn_tp_size + local_attn_tp_rank = local_tp_rank % local_attn_tp_size + + return local_attn_tp_rank, local_attn_tp_size, local_attn_dp_rank + + +def initialize_dp_attention( + server_args: ServerArgs, + model_config: ModelConfig, +): + global _ATTN_TP_GROUP, _ATTN_TP_RANK, _ATTN_TP_SIZE, _ATTN_DP_RANK, _ATTN_DP_SIZE + global _LOCAL_ATTN_DP_SIZE, _LOCAL_ATTN_DP_RANK, _ENABLE_DP_ATTENTION_FLAG + + from sglang.srt.layers.sampler import SYNC_TOKEN_IDS_ACROSS_TP + + enable_dp_attention = server_args.enable_dp_attention + tp_size = server_args.tp_size + dp_size = server_args.dp_size + moe_dense_tp_size = server_args.moe_dense_tp_size + pp_size = server_args.pp_size + + tp_rank = get_tensor_model_parallel_rank() + + _ENABLE_DP_ATTENTION_FLAG = enable_dp_attention + + _ATTN_TP_RANK, _ATTN_TP_SIZE, _ATTN_DP_RANK = compute_dp_attention_world_info( + enable_dp_attention, tp_rank, tp_size, dp_size + ) + _, _, _LOCAL_ATTN_DP_RANK = compute_dp_attention_local_info( + enable_dp_attention, tp_rank, tp_size, dp_size, moe_dense_tp_size + ) + + if enable_dp_attention: + _ATTN_DP_SIZE = dp_size + if moe_dense_tp_size is None: + _LOCAL_ATTN_DP_SIZE = _ATTN_DP_SIZE + else: + _LOCAL_ATTN_DP_SIZE = max(1, dp_size // (tp_size // moe_dense_tp_size)) + else: + _ATTN_DP_SIZE = 1 + _LOCAL_ATTN_DP_SIZE = 1 + + tp_group = get_tp_group() + _ATTN_TP_GROUP = GroupCoordinator( + [ + list(range(head, head + _ATTN_TP_SIZE)) + for head in range(0, pp_size * tp_size, _ATTN_TP_SIZE) + ], + tp_group.local_rank, + torch.distributed.get_backend(tp_group.device_group), + use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP, + use_pymscclpp=False, + use_custom_allreduce=False, + use_hpu_communicator=False, + use_xpu_communicator=False, + use_npu_communicator=False, + group_name="attention_tp", + ) + + _DpGatheredBufferWrapper.set_metadata( + hidden_size=model_config.hidden_size, + dtype=model_config.dtype, + device=torch.device(server_args.device), + ) + + +def is_dp_attention_enabled() -> bool: + return _ENABLE_DP_ATTENTION_FLAG + + +def get_attention_tp_group() -> GroupCoordinator: + assert _ATTN_TP_GROUP is not None, "dp attention not initialized!" + return _ATTN_TP_GROUP + + +def get_attention_tp_rank() -> int: + assert _ATTN_TP_RANK is not None, "dp attention not initialized!" + return _ATTN_TP_RANK + + +def get_attention_tp_size() -> int: + assert _ATTN_TP_SIZE is not None, "dp attention not initialized!" + return _ATTN_TP_SIZE + + +def get_attention_dp_rank() -> int: + assert _ATTN_DP_RANK is not None, "dp attention not initialized!" + return _ATTN_DP_RANK + + +def get_attention_dp_size() -> int: + assert _ATTN_DP_SIZE is not None, "dp attention not initialized!" + return _ATTN_DP_SIZE + + +def get_local_attention_dp_rank() -> int: + assert _LOCAL_ATTN_DP_RANK is not None, "dp attention not initialized!" + return _LOCAL_ATTN_DP_RANK + + +def get_local_attention_dp_size() -> int: + assert _LOCAL_ATTN_DP_SIZE is not None, "dp attention not initialized!" + return _LOCAL_ATTN_DP_SIZE + + +@contextmanager +def disable_dp_size(): + """Patch the tp group temporarily until this function ends. + + This method is for draft workers of speculative decoding to run draft model + with different tp degree from that of target model workers. + + Args: + tp_group (GroupCoordinator): the tp group coordinator + """ + global _ATTN_DP_SIZE + assert _ATTN_DP_SIZE is not None, "dp attention not initialized!" + + old_dp_size = _ATTN_DP_SIZE + _ATTN_DP_SIZE = 1 + try: + yield + finally: + _ATTN_DP_SIZE = old_dp_size + + +def get_dp_local_info(forward_batch: ForwardBatch) -> Tuple[torch.Tensor, torch.Tensor]: + # `get_dp_local_info` is only called in global DP gather and scatter. We use global DP rank here. + dp_rank = get_attention_dp_rank() + + if forward_batch.dp_local_start_pos is None: + cumtokens = torch.cumsum(forward_batch.global_num_tokens_gpu, dim=0) + if dp_rank == 0: + local_start_pos = torch.zeros_like(cumtokens[0]) + else: + local_start_pos = cumtokens[dp_rank - 1] + local_num_tokens = forward_batch.global_num_tokens_gpu[dp_rank] + + forward_batch.dp_local_start_pos = local_start_pos + forward_batch.dp_local_num_tokens = local_num_tokens + + return forward_batch.dp_local_start_pos, forward_batch.dp_local_num_tokens + + +@triton.jit +def memcpy_triton_kernel( + dst_ptr, + src_ptr, + offset_ptr, + sz_ptr, + offset_src: tl.constexpr, + chunk_size, # multiplied for offset and sz + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0).to(tl.int64) + offset = tl.load(offset_ptr).to(tl.int64) * chunk_size + sz = tl.load(sz_ptr).to(tl.int64) * chunk_size + + start_index = pid * BLOCK_SIZE + offs = tl.arange(0, BLOCK_SIZE) + mask = start_index + offs < sz + + if offset_src: + data = tl.load(src_ptr + offset + start_index + offs, mask=mask) + tl.store(dst_ptr + start_index + offs, data, mask=mask) + else: + data = tl.load(src_ptr + start_index + offs, mask=mask) + tl.store(dst_ptr + offset + start_index + offs, data, mask=mask) + + +def prod(x): + return functools.reduce(lambda a, b: a * b, x, 1) + + +def memcpy_triton(dst, src, dim, offset, sz, offset_src): + max_size = min(src.numel(), dst.numel()) + assert dim == 0, "dim != 0 unsupported" + assert src.shape[1:] == dst.shape[1:], "src and dst must have same shape" + chunk_size = prod(src.shape[1:]) + BLOCK_SIZE = 8192 + grid = (triton.cdiv(max_size, BLOCK_SIZE),) + + memcpy_triton_kernel[grid](dst, src, offset, sz, offset_src, chunk_size, BLOCK_SIZE) + + +def _dp_gather_via_all_reduce( + global_tokens: torch.Tensor, + local_tokens: torch.Tensor, + forward_batch: ForwardBatch, + is_partial: bool, +): + local_start_pos, local_num_tokens = get_dp_local_info(forward_batch) + + global_tokens.fill_(0) + assert local_tokens.is_contiguous() + assert global_tokens.is_contiguous() + + if local_tokens.shape[0] > 0 and (is_partial or get_attention_tp_rank() == 0): + assert ( + local_tokens.untyped_storage() is not global_tokens.untyped_storage() + ), "aliasing between global_tokens and local_tokens not allowed" + + memcpy_triton( + global_tokens, local_tokens, 0, local_start_pos, local_num_tokens, False + ) + + # Input IDs are in int 32. We should use inplace_all_reduce for local case because of custom all reduce. + NUM_GPUS_PER_NODE = 8 + if ( + not local_tokens.dtype.is_floating_point + and get_tensor_model_parallel_world_size() <= NUM_GPUS_PER_NODE + ): + torch.ops.sglang.inplace_all_reduce( + global_tokens, group_name=get_tp_group().unique_name + ) + + else: + global_tokens[:] = tensor_model_parallel_all_reduce(global_tokens) + + +def _dp_gather_via_all_gather( + global_tokens: torch.Tensor, + local_tokens: torch.Tensor, + forward_batch: ForwardBatch, + is_partial: bool, +): + if get_attention_tp_size() == 1: + get_tp_group().all_gather_into_tensor(global_tokens, local_tokens) + return + + if not is_partial: + if get_attention_tp_rank() != 0: + local_tokens.fill_(0) + scattered_local_tokens = local_tokens.tensor_split(get_attention_tp_size())[ + get_attention_tp_rank() + ] + get_attention_tp_group().reduce_scatter_tensor(scattered_local_tokens, local_tokens) + get_tp_group().all_gather_into_tensor(global_tokens, scattered_local_tokens) + + +def _dp_gather( + global_tokens: torch.Tensor, + local_tokens: torch.Tensor, + forward_batch: ForwardBatch, + is_partial: bool, +): + if forward_batch.dp_padding_mode.is_max_len(): + _dp_gather_via_all_gather( + global_tokens, local_tokens, forward_batch, is_partial + ) + else: + _dp_gather_via_all_reduce( + global_tokens, local_tokens, forward_batch, is_partial + ) + + +def dp_gather_partial( + global_tokens: torch.Tensor, + local_tokens: torch.Tensor, + forward_batch: ForwardBatch, +): + _dp_gather(global_tokens, local_tokens, forward_batch, is_partial=True) + + +def dp_gather_replicate( + global_tokens: torch.Tensor, + local_tokens: torch.Tensor, + forward_batch: ForwardBatch, +): + _dp_gather(global_tokens, local_tokens, forward_batch, is_partial=False) + + +def dp_scatter( + local_tokens: torch.Tensor, # output + global_tokens: torch.Tensor, # input + forward_batch: ForwardBatch, +): + # local_num_tokens is not necessarily the same as local_tokens.shape[0], + # since local_tokens may be padded for cuda graph + local_start_pos, local_num_tokens = get_dp_local_info(forward_batch) + + local_tokens.fill_(0) + assert local_tokens.is_contiguous() + assert global_tokens.is_contiguous() + if local_tokens.shape[0] > 0: + assert ( + local_tokens.untyped_storage() is not global_tokens.untyped_storage() + ), "aliasing between local_tokens and global_tokens not allowed" + + memcpy_triton( + local_tokens, global_tokens, 0, local_start_pos, local_num_tokens, True + ) + + +def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor): + if get_tensor_model_parallel_world_size() == get_attention_dp_size(): + get_tp_group().reduce_scatter_tensor(output, input) + else: + scattered_local_tokens = input.tensor_split( + get_tensor_model_parallel_world_size() + )[get_tensor_model_parallel_rank()] + get_tp_group().reduce_scatter_tensor(scattered_local_tokens, input) + get_attention_tp_group().all_gather_into_tensor(output, scattered_local_tokens) + + +def attn_tp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor): + return get_attention_tp_group().reduce_scatter_tensor(output, input) + + +def attn_tp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor): + return get_attention_tp_group().all_gather_into_tensor(output, input) + + +def attn_tp_all_gather(output_list: List[torch.Tensor], input: torch.Tensor): + return get_attention_tp_group().all_gather(input, output_tensor_list=output_list) diff --git a/python/sglang/srt/layers/elementwise.py b/python/sglang/srt/layers/elementwise.py new file mode 100644 index 000000000..e05d88b32 --- /dev/null +++ b/python/sglang/srt/layers/elementwise.py @@ -0,0 +1,582 @@ +from typing import Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.utils import is_hip + +_is_hip = is_hip() + + +fused_softcap_autotune = triton.autotune( + configs=[ + triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=4), + triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=4), + triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=4), + triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 32768}, num_warps=32), + ], + key=["n_ele"], +) + + +@triton.jit +def fused_softcap_kernel( + output_ptr, + input_ptr, + n_ele, + softcap_const: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_ele + x = tl.load(input_ptr + offsets, mask=mask) + fx = x.to(tl.float32) + fxs = fx / softcap_const + exped = tl.exp(2 * fxs) + top = exped - 1 + bottom = exped + 1 + output = top / bottom * softcap_const + tl.store(output_ptr + offsets, output, mask=mask) + + +fused_softcap_kernel_autotuned = fused_softcap_autotune(fused_softcap_kernel) + + +def fused_softcap(x, softcap_const, autotune=False): + output = torch.empty_like(x, dtype=torch.float32) + n_elements = output.numel() + if autotune: + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + fused_softcap_kernel_autotuned[grid](output, x, n_elements, softcap_const) + else: + fused_softcap_kernel[(triton.cdiv(n_elements, 128),)]( + output, x, n_elements, softcap_const, BLOCK_SIZE=128, num_warps=8 + ) + return output + + +# cast to float + softcap +class Softcap: + def __init__(self, softcap_const: float): + self.softcap_const = softcap_const + + def __call__(self, *args, **kwargs): + return self.forward(*args, **kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.is_cuda: + return self.forward_cuda(x) + else: + return self.forward_native(x) + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + return torch.tanh(x.float() / self.softcap_const) * self.softcap_const + + def forward_cuda(self, x: torch.Tensor, autotune=False) -> torch.Tensor: + return fused_softcap(x, self.softcap_const, autotune=autotune) + + +rmsnorm_autotune = triton.autotune( + configs=[ + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4, num_stages=1), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8, num_stages=1), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16, num_stages=1), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8, num_stages=8), + triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16, num_stages=8), + triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=8, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=16, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=8, num_stages=1), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=16, num_stages=1), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32, num_stages=1), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=8, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=16, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=8), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=16), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=8, num_stages=1), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=16, num_stages=1), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32, num_stages=1), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=8, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=16, num_stages=4), + triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32, num_stages=4), + ], + key=["hidden_dim"], +) + + +@triton.jit +def fused_dual_residual_rmsnorm_kernel( + output_ptr, + mid_ptr, + activ_ptr, + residual_ptr, + weight1_ptr, + weight2_ptr, + eps: tl.constexpr, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + input_start = pid * hidden_dim + + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_dim + + a_ = tl.load(activ_ptr + input_start + offsets, mask=mask, other=0.0) + a = a_.to(tl.float32) + rms = tl.sqrt(tl.sum(a * a, axis=0) / hidden_dim + eps) + + r = tl.load(residual_ptr + input_start + offsets, mask=mask, other=0.0) + w1_ = tl.load(weight1_ptr + offsets, mask=mask, other=0.0) + w1 = w1_.to(tl.float32) + + a2r = r + (a / rms * w1).to(r.dtype) + tl.store( + mid_ptr + input_start + offsets, + a2r, + mask=mask, + ) + + a2r = a2r.to(tl.float32) + rms2 = tl.sqrt(tl.sum(a2r * a2r, axis=0) / hidden_dim + eps) + + w2_ = tl.load(weight2_ptr + offsets, mask=mask, other=0.0) + w2 = w2_.to(tl.float32) + + tl.store( + output_ptr + input_start + offsets, + a2r / rms2 * w2, # implicitly casts to output dtype here + mask=mask, + ) + + +fused_dual_residual_rmsnorm_kernel_autotune = rmsnorm_autotune( + fused_dual_residual_rmsnorm_kernel +) + + +def fused_dual_residual_rmsnorm(x, residual, weight1, weight2, eps, autotune=False): + assert len(x.shape) == 2 + assert x.shape == residual.shape and x.dtype == residual.dtype + output, mid = torch.empty_like(x), torch.empty_like(x) + bs, hidden_dim = x.shape + if autotune: + fused_dual_residual_rmsnorm_kernel_autotune[(bs,)]( + output, mid, x, residual, weight1, weight2, eps=eps, hidden_dim=hidden_dim + ) + else: + max_warps = 16 if _is_hip else 32 + config = { + "BLOCK_SIZE": triton.next_power_of_2(hidden_dim), + "num_warps": max( + min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), max_warps), 4 + ), + } + + fused_dual_residual_rmsnorm_kernel[(bs,)]( + output, + mid, + x, + residual, + weight1, + weight2, + eps=eps, + hidden_dim=hidden_dim, + **config, + ) + + return output, mid + + +@triton.jit +def fused_rmsnorm_kernel( + output_ptr, + activ_ptr, + weight_ptr, + eps: tl.constexpr, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + input_start = pid * hidden_dim + + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_dim + + a_ = tl.load(activ_ptr + input_start + offsets, mask=mask, other=0.0) + a = a_.to(tl.float32) + rms = tl.sqrt(tl.sum(a * a, axis=0) / hidden_dim + eps) + + w1_ = tl.load(weight_ptr + offsets, mask=mask, other=0.0) + w1 = w1_.to(tl.float32) + + a_rms = a / rms * w1 + + tl.store( + output_ptr + input_start + offsets, + a_rms, # implicitly casts to output dtype here + mask=mask, + ) + + +def fused_rmsnorm(x, weight, eps, autotune=False, inplace=False): + assert len(x.shape) == 2 + if inplace: + output = x + else: + output = torch.empty_like(x) + bs, hidden_dim = x.shape + max_warps = 16 if _is_hip else 32 + config = { + "BLOCK_SIZE": triton.next_power_of_2(hidden_dim), + "num_warps": max( + min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), max_warps), 4 + ), + } + + fused_rmsnorm_kernel[(bs,)]( + output, x, weight, eps=eps, hidden_dim=hidden_dim, **config + ) + return output + + +class FusedDualResidualRMSNorm: + """ + Fused implementation of + y = RMSNorm2(RMSNorm1(x) + residual)) + """ + + def __init__(self, rmsnorm1, rmsnorm2) -> None: # the one after rmsnorm1 + self.rmsnorm1 = rmsnorm1 + self.rmsnorm2 = rmsnorm2 + self.variance_epsilon = self.rmsnorm1.variance_epsilon + assert self.rmsnorm1.variance_epsilon == self.rmsnorm2.variance_epsilon + assert self.rmsnorm1.weight.shape == self.rmsnorm2.weight.shape + + def __call__(self, *args, **kwargs): + return self.forward(*args, **kwargs) + + def forward( + self, x: torch.Tensor, residual: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + if x.is_cuda: + return self.forward_cuda(x, residual) + else: + return self.forward_flashinfer(x, residual) + + def forward_cuda( + self, x: torch.Tensor, residual: torch.Tensor, autotune=False + ) -> Tuple[torch.Tensor, torch.Tensor]: + return fused_dual_residual_rmsnorm( + x, + residual, + self.rmsnorm1.weight, + self.rmsnorm2.weight, + self.variance_epsilon, + autotune=autotune, + ) + + def forward_flashinfer( + self, + x: torch.Tensor, + residual: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + normed1 = self.rmsnorm1(x) + residual = normed1 + residual + return self.rmsnorm2(residual), residual + + def forward_native( + self, + x: torch.Tensor, + residual: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + normed1 = self.rmsnorm1.forward_native(x) + residual = normed1 + residual + return self.rmsnorm2.forward_native(residual), residual + + +@triton.jit +def experts_combine_kernel( + out_hidden_states, + moe_hidden_states, + mlp_hidden_states, + combine_k: tl.constexpr, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + start_index_mlp = pid * hidden_dim + start_index_rmoe = pid * hidden_dim * combine_k + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_dim + combine_k_offsets = tl.arange(0, combine_k) + + moe_x = tl.load( + moe_hidden_states + + start_index_rmoe + + combine_k_offsets[:, None] * hidden_dim + + offsets[None, :], + mask=mask[None, :], + other=0.0, + ) + moe_x = tl.sum(moe_x, axis=0) + mlp_x = tl.load(mlp_hidden_states + start_index_mlp + offsets, mask=mask, other=0.0) + combined_x = (moe_x + mlp_x) / 1.4142135623730951 + + tl.store(out_hidden_states + start_index_mlp + offsets, combined_x, mask=mask) + + +def experts_combine_triton(moe_hidden_states, mlp_hidden_states, output_buffer=None): + assert moe_hidden_states.is_contiguous() + assert mlp_hidden_states.is_contiguous() + + if len(moe_hidden_states.shape) == 2: + combine_k = 1 # pre-combined + else: + combine_k = moe_hidden_states.shape[1] + + if output_buffer is None: + out_hidden_states = torch.empty_like(mlp_hidden_states) + else: + flat_output_buffer = output_buffer.view(mlp_hidden_states.dtype).reshape(-1) + assert flat_output_buffer.numel() >= mlp_hidden_states.numel() + out_hidden_states = flat_output_buffer[: mlp_hidden_states.numel()].reshape( + mlp_hidden_states.shape + ) + + bs, hidden_dim = mlp_hidden_states.shape + + config = { + "BLOCK_SIZE": triton.next_power_of_2(hidden_dim), + "num_warps": max( + min(triton.next_power_of_2(triton.cdiv(hidden_dim, 1024)), 8), 4 + ), + } + + experts_combine_kernel[(bs,)]( + out_hidden_states, + moe_hidden_states, + mlp_hidden_states, + combine_k, + hidden_dim, + **config, + ) + return out_hidden_states + + +# gelu on first half of vector +@triton.jit +def gelu_and_mul_kernel( + out_hidden_states_ptr, # (bs, hidden_dim) + out_scales_ptr, # (bs,) + hidden_states_ptr, # (bs, hidden_dim * 2) + quant_max: tl.constexpr, + static_scale: tl.constexpr, + hidden_dim: tl.constexpr, # the output hidden_dim + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + + input_start = pid * hidden_dim * 2 + output_start = pid * hidden_dim + + input1_offs = tl.arange(0, BLOCK_SIZE) + mask = tl.arange(0, BLOCK_SIZE) < hidden_dim # shared for input1, input3, output + input3_offs = hidden_dim + tl.arange(0, BLOCK_SIZE) + output_offs = tl.arange(0, BLOCK_SIZE) + + x1 = tl.load( + hidden_states_ptr + input_start + input1_offs, mask=mask, other=0.0 + ).to(tl.float32) + x3 = tl.load( + hidden_states_ptr + input_start + input3_offs, mask=mask, other=0.0 + ).to(tl.float32) + + # gelu + # cast down before mul to better match training? + gelu_x1 = 0.5 * (1.0 + tl.erf(x1 * 0.7071067811865475)) * x1 + out = x3 * gelu_x1.to(hidden_states_ptr.dtype.element_ty) + + if quant_max is not None: + raise NotImplementedError() + + tl.store(out_hidden_states_ptr + output_start + output_offs, out, mask=mask) + + +def gelu_and_mul_triton( + hidden_states, + scales=None, + quantize=None, # dtype to quantize to + out=None, +): + bs, in_hidden_dim = hidden_states.shape + hidden_dim = in_hidden_dim // 2 + + if out is None: + out_hidden_states = torch.empty( + (bs, hidden_dim), + dtype=quantize or hidden_states.dtype, + device=hidden_states.device, + ) + else: + assert out.shape == (bs, hidden_dim) + assert out.dtype == (quantize or hidden_states.dtype) + out_hidden_states = out + out_scales = None + static_scale = False + if quantize is not None: + if scales is None: + out_scales = torch.empty( + (bs,), dtype=torch.float32, device=hidden_states.device + ) + else: + out_scales = scales + static_scale = True + + max_warps = 16 if _is_hip else 32 + config = { + # 8 ele per thread (not tuned) + "num_warps": max( + min(triton.next_power_of_2(triton.cdiv(hidden_dim, 8 * 32)), max_warps), 4 + ), + } + + gelu_and_mul_kernel[(bs,)]( + out_hidden_states, + out_scales, + hidden_states, + quant_max=torch.finfo(quantize).max if quantize is not None else None, + static_scale=static_scale, + hidden_dim=hidden_dim, + BLOCK_SIZE=triton.next_power_of_2(hidden_dim), + **config, + ) + + if quantize is not None: + return out_hidden_states, out_scales + else: + return out_hidden_states, None + + +# silu on first half of vector +@triton.jit +def silu_and_mul_kernel( + out_hidden_states_ptr, # (bs, hidden_dim) + out_scales_ptr, # (bs,) + hidden_states_ptr, # (bs, hidden_dim * 2) + quant_max: tl.constexpr, + static_scale: tl.constexpr, + hidden_dim: tl.constexpr, # the output hidden_dim + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + + input_start = pid * hidden_dim * 2 + output_start = pid * hidden_dim + + input1_offs = tl.arange(0, BLOCK_SIZE) + mask = tl.arange(0, BLOCK_SIZE) < hidden_dim # shared for input1, input3, output + input3_offs = hidden_dim + tl.arange(0, BLOCK_SIZE) + output_offs = tl.arange(0, BLOCK_SIZE) + + x1 = tl.load( + hidden_states_ptr + input_start + input1_offs, mask=mask, other=0.0 + ).to(tl.float32) + x3 = tl.load( + hidden_states_ptr + input_start + input3_offs, mask=mask, other=0.0 + ).to(tl.float32) + + # silu + # cast down before mul to better match training? + silu_x1 = x1 * tl.sigmoid(x1) + out = x3 * silu_x1.to(hidden_states_ptr.dtype.element_ty) + + if quant_max is not None: + raise NotImplementedError() + + tl.store(out_hidden_states_ptr + output_start + output_offs, out, mask=mask) + + +def silu_and_mul_triton( + hidden_states, + scales=None, + quantize=None, # dtype to quantize to + out=None, +): + bs, in_hidden_dim = hidden_states.shape + hidden_dim = in_hidden_dim // 2 + + if out is None: + out_hidden_states = torch.empty( + (bs, hidden_dim), + dtype=quantize or hidden_states.dtype, + device=hidden_states.device, + ) + else: + assert out.shape == (bs, hidden_dim) + assert out.dtype == (quantize or hidden_states.dtype) + out_hidden_states = out + out_scales = None + static_scale = False + if quantize is not None: + if scales is None: + out_scales = torch.empty( + (bs,), dtype=torch.float32, device=hidden_states.device + ) + else: + out_scales = scales + static_scale = True + + max_warps = 16 if _is_hip else 32 + config = { + # 8 ele per thread (not tuned) + "num_warps": max( + min(triton.next_power_of_2(triton.cdiv(hidden_dim, 8 * 32)), max_warps), 4 + ), + } + + silu_and_mul_kernel[(bs,)]( + out_hidden_states, + out_scales, + hidden_states, + quant_max=torch.finfo(quantize).max if quantize is not None else None, + static_scale=static_scale, + hidden_dim=hidden_dim, + BLOCK_SIZE=triton.next_power_of_2(hidden_dim), + **config, + ) + + if quantize is not None: + return out_hidden_states, out_scales + else: + return out_hidden_states, None diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py new file mode 100644 index 000000000..81280db0a --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_comm_fusion.py @@ -0,0 +1,230 @@ +import logging +from typing import Optional, Tuple + +import torch +import torch.distributed as dist + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.utils import ( + direct_register_custom_op, + is_flashinfer_available, + supports_custom_op, +) + +logger = logging.getLogger(__name__) + +_flashinfer_comm = None +_workspace_manager = None + +if is_flashinfer_available(): + try: + import flashinfer.comm as comm + + _flashinfer_comm = comm + except ImportError: + logger.warning( + "flashinfer.comm is not available, falling back to standard " + "implementation" + ) + + +class FlashInferWorkspaceManager: + def __init__(self): + self.workspace_tensor = None + self.ipc_handles = None + self.world_size = None + self.rank = None + self.initialized = False + + def initialize( + self, + world_size: int, + rank: int, + max_token_num: int, + hidden_dim: int, + group=None, + use_fp32_lamport: bool = False, + ): + """Initialize workspace""" + if self.initialized and self.world_size == world_size: + return + + if _flashinfer_comm is None: + logger.warning( + "FlashInfer comm not available, skipping workspace " "initialization" + ) + return + + self.cleanup() + + self.ipc_handles, self.workspace_tensor = ( + comm.trtllm_create_ipc_workspace_for_all_reduce_fusion( + rank, + world_size, + max_token_num, + hidden_dim, + group=group, + use_fp32_lamport=use_fp32_lamport, + ) + ) + + self.world_size = world_size + self.rank = rank + self.initialized = True + + logger.info( + f"FlashInfer workspace initialized for rank {rank}, " + f"world_size {world_size}" + ) + + def cleanup(self): + """Clean up workspace""" + if self.initialized and self.ipc_handles is not None: + try: + _flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce( + self.ipc_handles, group=dist.group.WORLD + ) + except Exception as e: + logger.warning(f"Failed to cleanup FlashInfer workspace: {e}") + finally: + self.workspace_tensor = None + self.ipc_handles = None + self.initialized = False + + +_workspace_manager = FlashInferWorkspaceManager() + + +def ensure_workspace_initialized( + max_token_num: int = 2048, hidden_dim: int = 4096, use_fp32_lamport: bool = False +): + """Ensure workspace is initialized""" + if not is_flashinfer_available() or _flashinfer_comm is None: + return False + + world_size = get_tensor_model_parallel_world_size() + if world_size <= 1: + return False + + rank = dist.get_rank() + + if ( + not _workspace_manager.initialized + or _workspace_manager.world_size != world_size + ): + _workspace_manager.initialize( + world_size=world_size, + rank=rank, + max_token_num=max_token_num, + hidden_dim=hidden_dim, + use_fp32_lamport=use_fp32_lamport, + ) + + return _workspace_manager.initialized + + +def flashinfer_allreduce_residual_rmsnorm( + input_tensor: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + max_token_num: int = 2048, + use_oneshot: Optional[bool] = None, + trigger_completion_at_end: bool = False, + fp32_acc: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Use FlashInfer's fused allreduce + residual + RMS norm operation + + Args: + input_tensor: Input tensor that needs allreduce + residual: Residual tensor + weight: RMS norm weight + eps: RMS norm epsilon + max_token_num: Maximum token number + use_oneshot: Whether to use oneshot mode + trigger_completion_at_end: Whether to trigger completion at end + fp32_acc: Whether to use fp32 precision + + Returns: + Tuple[torch.Tensor, torch.Tensor]: (norm_output, residual_output) + """ + if not is_flashinfer_available() or _flashinfer_comm is None: + logger.debug( + "FlashInfer not available, falling back to standard " "implementation" + ) + return None, None + + world_size = get_tensor_model_parallel_world_size() + if world_size <= 1: + logger.debug("Single GPU, no need for allreduce fusion") + return None, None + + if not ensure_workspace_initialized( + max_token_num=max_token_num, + hidden_dim=input_tensor.shape[-1], + use_fp32_lamport=(input_tensor.dtype == torch.float32), + ): + logger.debug("FlashInfer workspace not available") + return None, None + + token_num, hidden_dim = input_tensor.shape + + residual_out = torch.empty_like(residual) + norm_out = torch.empty_like(input_tensor) + + _flashinfer_comm.trtllm_allreduce_fusion( + allreduce_in=input_tensor, + world_size=world_size, + world_rank=dist.get_rank(), + token_num=token_num, + hidden_dim=hidden_dim, + workspace_ptrs=_workspace_manager.workspace_tensor, + launch_with_pdl=True, + use_oneshot=use_oneshot, + trigger_completion_at_end=trigger_completion_at_end, + fp32_acc=fp32_acc, + pattern_code=(_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm), + allreduce_out=None, + residual_in=residual, + residual_out=residual_out, + norm_out=norm_out, + quant_out=None, + scale_out=None, + rms_gamma=weight, + rms_eps=eps, + scale_factor=None, + layout_code=None, + ) + + return norm_out, residual_out + + +def fake_flashinfer_allreduce_residual_rmsnorm( + input_tensor: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + max_token_num: int = 2048, + use_oneshot: Optional[bool] = None, + trigger_completion_at_end: bool = False, + fp32_acc: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + residual_out = torch.empty_like(residual) + norm_out = torch.empty_like(input_tensor) + return norm_out, residual_out + + +if supports_custom_op(): + direct_register_custom_op( + "flashinfer_allreduce_residual_rmsnorm", + flashinfer_allreduce_residual_rmsnorm, + mutates_args=["input_tensor", "residual", "weight"], + fake_impl=fake_flashinfer_allreduce_residual_rmsnorm, + ) + + +def cleanup_flashinfer_workspace(): + global _workspace_manager + if _workspace_manager is not None: + _workspace_manager.cleanup() diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py new file mode 100644 index 000000000..12cd49adf --- /dev/null +++ b/python/sglang/srt/layers/layernorm.py @@ -0,0 +1,343 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Fused operators for normalization layers.""" + +import logging +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +from packaging.version import Version + +from sglang.srt.custom_op import CustomOp +from sglang.srt.utils import ( + cpu_has_amx_support, + get_bool_env_var, + is_cpu, + is_cuda, + is_flashinfer_available, + is_hip, + is_npu, + is_xpu, + supports_custom_op, +) + +_is_cuda = is_cuda() +_is_flashinfer_available = is_flashinfer_available() +_is_hip = is_hip() +_is_npu = is_npu() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() +_is_xpu = is_xpu() + +if _is_cuda: + if _is_flashinfer_available: + from flashinfer.norm import fused_add_rmsnorm + else: + from sgl_kernel import fused_add_rmsnorm + from sgl_kernel import gemma_fused_add_rmsnorm, gemma_rmsnorm, rmsnorm + +if _use_aiter: + from aiter import rmsnorm2d_fwd as rms_norm + from aiter import rmsnorm2d_fwd_with_add as fused_add_rms_norm +elif _is_hip: + import vllm + from vllm._custom_ops import fused_add_rms_norm, rms_norm + + _vllm_version = Version(vllm.__version__) + +logger = logging.getLogger(__name__) + +if _is_npu: + import torch_npu + + +class RMSNorm(CustomOp): + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + var_hidden_size: Optional[int] = None, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.hidden_size = hidden_size + self.variance_size_override = ( + None if var_hidden_size == hidden_size else var_hidden_size + ) + if _use_aiter: + self._forward_method = self.forward_aiter + + def forward_cuda( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if self.variance_size_override is not None: + return self.forward_native(x, residual) + if residual is not None: + fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon) + return x, residual + out = rmsnorm(x, self.weight.data, self.variance_epsilon) + return out + + def forward_npu( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if residual is not None: + out, _, residual_out = torch_npu.npu_add_rms_norm( + residual, x, self.weight.data, self.variance_epsilon + ) + return out, residual_out + return torch_npu.npu_rms_norm(x, self.weight.data, self.variance_epsilon)[0] + + def forward_aiter( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if residual is not None: + residual_out = torch.empty_like(x) + output = torch.empty_like(x) + fused_add_rms_norm( + output, + x, + residual, + residual_out, + self.weight.data, + self.variance_epsilon, + ) + return output, residual_out + return rms_norm(x, self.weight.data, self.variance_epsilon) + + def forward_hip( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ): + if not x.is_contiguous(): + x = x.contiguous() + + if residual is not None: + try: + output = torch.empty_like(x) + residual_out = torch.empty_like(x) + fused_add_rms_norm( + output, + x, + residual_out, + residual, + self.weight.data, + self.variance_epsilon, + ) + return output, residual_out + except TypeError: + fused_add_rms_norm( + x, + residual, + self.weight.data, + self.variance_epsilon, + ) + return x, residual + + out = torch.empty_like(x) + rms_norm(out, x, self.weight.data, self.variance_epsilon) + return out + + + def forward_native( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if not x.is_contiguous(): + x = x.contiguous() + orig_dtype = x.dtype + x = x.to(torch.float32) + if residual is not None: + x = x + residual.to(torch.float32) + residual = x.to(orig_dtype) + + hidden_size = x.shape[-1] + if hidden_size != self.hidden_size: + raise ValueError( + "Expected hidden_size to be " + f"{self.hidden_size}, but found: {hidden_size}" + ) + + if self.variance_size_override is None: + x_var = x + else: + if hidden_size < self.variance_size_override: + raise ValueError( + "Expected hidden_size to be at least " + f"{self.variance_size_override}, but found: {hidden_size}" + ) + + x_var = x[..., : self.variance_size_override] + + variance = x_var.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + self.variance_epsilon) + x = (x * self.weight).to(orig_dtype) + if residual is None: + return x + else: + return x, residual + + def forward_cpu( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if _is_cpu_amx_available: + if residual is not None: + torch.ops.sgl_kernel.fused_add_rmsnorm_cpu( + x, residual, self.weight.data, self.variance_epsilon + ) + return x, residual + return torch.ops.sgl_kernel.rmsnorm_cpu( + x, self.weight.data, self.variance_epsilon + ) + else: + return self.forward_native(x, residual) + + def forward_with_allreduce_fusion( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + Forward method with allreduce fusion, prioritizing flashinfer fused operations + """ + if residual is not None: + from sglang.srt.distributed import get_tensor_model_parallel_world_size + from sglang.srt.layers.flashinfer_comm_fusion import ( + flashinfer_allreduce_residual_rmsnorm, + ) + + fused_op = ( + torch.ops.sglang.flashinfer_allreduce_residual_rmsnorm + if supports_custom_op() + else flashinfer_allreduce_residual_rmsnorm + ) + + if get_tensor_model_parallel_world_size() > 1: + fused_result = fused_op( + input_tensor=x, + residual=residual, + weight=self.weight, + eps=self.variance_epsilon, + ) + if fused_result[0] is not None: + return fused_result + + return self.forward(x, residual) + + +class GemmaRMSNorm(CustomOp): + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + self.variance_epsilon = eps + + # Re-dispatch + if _is_hip: + self._forward_method = self.forward_native + + def forward_native( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + orig_dtype = x.dtype + if residual is not None: + x = x + residual + residual = x + + x = x.float() + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + self.variance_epsilon) + x = x * (1.0 + self.weight.float()) + x = x.to(orig_dtype) + return x if residual is None else (x, residual) + + def forward_cuda( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if residual is not None: + gemma_fused_add_rmsnorm( + x, residual, self.weight.data, self.variance_epsilon + ) + return x, residual + out = gemma_rmsnorm(x, self.weight.data, self.variance_epsilon) + return out + + def forward_npu( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if residual is not None: + x = x + residual + residual = x + + x, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.variance_epsilon) + return x if residual is None else (x, residual) + + +class Gemma3RMSNorm(CustomOp): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + # Re-dispatch + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward_native(self, x): + output = self._norm(x.float()) + # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16) + # See https://github.com/huggingface/transformers/pull/29402 + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + def forward_cuda(self, x): + return self.forward_native(x) + + def forward_npu(self, x): + output, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.eps) + return output + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.eps}" + + +if not ( + _is_cuda or _is_hip or _is_npu or (_is_cpu and _is_cpu_amx_available) or _is_xpu +): + logger.info( + "sgl-kernel layernorm implementation is not available on current platform. Fallback to other kernel libraries." + ) + from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py new file mode 100644 index 000000000..df2b77e08 --- /dev/null +++ b/python/sglang/srt/layers/linear.py @@ -0,0 +1,1337 @@ +"""Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/layers/linear.py""" + +from __future__ import annotations + +import itertools +import logging +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +import torch +from torch.nn.parameter import Parameter, UninitializedParameter + +from sglang.srt.distributed import ( + divide, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + parallel_state, + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.layers.parameter import ( + BasevLLMParameter, + BlockQuantScaleParameter, + PackedColumnParameter, + PackedvLLMParameter, + PerTensorScaleParameter, + RowvLLMParameter, + _ColumnvLLMParameter, +) +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.utils import is_cpu, is_npu, set_weight_attrs + +if TYPE_CHECKING: + from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, + ) + +logger = logging.getLogger(__name__) + +WEIGHT_LOADER_V2_SUPPORTED = [ + "CompressedTensorsLinearMethod", + "AWQMarlinLinearMethod", + "AWQLinearMethod", + "GPTQMarlinLinearMethod", + "Fp8LinearMethod", + "BlockInt8LinearMethod", + "MarlinLinearMethod", + "QQQLinearMethod", + "GPTQMarlin24LinearMethod", + "TPUInt8LinearMethod", + "GPTQLinearMethod", + "FBGEMMFp8LinearMethod", + "ModelOptFp8LinearMethod", + "ModelOptFp4LinearMethod", + "IPEXAWQLinearMethod", + "PetitNvFp4LinearMethod", +] + +_is_cpu = is_cpu() +_is_npu = is_npu() + + +def adjust_marlin_shard(param, shard_size, shard_offset): + marlin_tile_size = getattr(param, "marlin_tile_size", None) + if marlin_tile_size is None: + return shard_size, shard_offset + + return shard_size * marlin_tile_size, shard_offset * marlin_tile_size + + +def adjust_bitsandbytes_4bit_shard( + param: Parameter, shard_offsets: Dict[str, Tuple[int, int]], loaded_shard_id: str +) -> Tuple[int, int]: + """Adjust the quantization offsets and sizes for BitsAndBytes sharding.""" + + total, _ = shard_offsets["total"] + orig_offset, orig_size = shard_offsets[loaded_shard_id] + + quantized_total = param.data.shape[0] + quantized_offset = orig_offset * quantized_total // total + quantized_size = orig_size * quantized_total // total + + return quantized_size, quantized_offset + + +def adjust_scalar_to_fused_array(param, loaded_weight, shard_id): + """For fused modules (QKV and MLP) we have an array of length + N that holds 1 scale for each "logical" matrix. So the param + is an array of length N. The loaded_weight corresponds to + one of the shards on disk. Here, we slice the param based on + the shard_id for loading. + """ + qkv_idxs = {"q": 0, "k": 1, "v": 2} + + if isinstance(shard_id, str): + shard_id = qkv_idxs[shard_id] + elif not isinstance(shard_id, int): + raise ValueError(f"Unknown Shard Id {shard_id}") + + # AutoFP8 scales do not have a shape + # compressed-tensors scales do have a shape + if len(loaded_weight.shape) != 0: + assert loaded_weight.shape[0] == 1 + loaded_weight = loaded_weight[0] + + return param[shard_id], loaded_weight + + +def adjust_shard_offsets(shard_offsets, loaded_weight, dim): + actual_weight_size = loaded_weight.size(dim) + target_weight_size = shard_offsets[-1][-1] + shard_offsets[-1][-2] + if actual_weight_size != target_weight_size: + new_shard_offsets = [] + new_offset = 0 + for shard_id, shard_offset, shard_size in shard_offsets: + actual_shard_size = actual_weight_size * shard_size // target_weight_size + new_shard_offsets.append((shard_id, new_offset, actual_shard_size)) + new_offset += actual_shard_size + return new_shard_offsets + return shard_offsets + + +class LinearBase(torch.nn.Module): + """Base linear layer. + + Args: + input_size: input dimension of the linear layer. + output_size: output dimension of the linear layer. + bias: If true, add bias. + skip_bias_add: If true, skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + """ + + def __init__( + self, + input_size: int, + output_size: int, + skip_bias_add: bool = False, + params_dtype: Optional[torch.dtype] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + + # Keep input parameters + self.input_size = input_size + self.output_size = output_size + self.skip_bias_add = skip_bias_add + if params_dtype is None: + params_dtype = torch.get_default_dtype() + self.params_dtype = params_dtype + if quant_config is None: + self.quant_method: Optional[QuantizeMethodBase] = UnquantizedLinearMethod() + else: + self.quant_method = quant_config.get_quant_method(self, prefix=prefix) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + +class ReplicatedLinear(LinearBase): + """Replicated linear layer. + + Args: + input_size: input dimension of the linear layer. + output_size: output dimension of the linear layer. + bias: If true, add bias. + skip_bias_add: If true, skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: Optional[torch.dtype] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__( + input_size, + output_size, + skip_bias_add, + params_dtype, + quant_config, + prefix=prefix, + ) + + # All the linear layer supports quant method. + assert self.quant_method is not None + self.quant_method.create_weights( + self, + self.input_size, + [self.output_size], + self.input_size, + self.output_size, + self.params_dtype, + weight_loader=self.weight_loader, + ) + + if bias: + self.bias = Parameter( + torch.empty(self.output_size, dtype=self.params_dtype) + ) + set_weight_attrs( + self.bias, + { + "output_dim": 0, + "weight_loader": self.weight_loader, + }, + ) + else: + self.register_parameter("bias", None) + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + # If the weight on disk does not have a shape, give it one + # (such scales for AutoFp8). + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + # The per-tensor quant-scale must be 1 dimension + if _is_npu: + if param.size() != loaded_weight.size() and param.size(0) == 1: + if torch.allclose(loaded_weight, loaded_weight[0]): + loaded_weight = loaded_weight[:1] + else: + raise ValueError(f"{loaded_weight} are not all equal") + + assert param.size() == loaded_weight.size() + param.data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + bias = self.bias if not self.skip_bias_add else None + assert self.quant_method is not None + output = self.quant_method.apply(self, x, bias) + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + def extra_repr(self) -> str: + s = f"in_features={self.input_size}" + s += f", output_features={self.output_size}" + s += f", bias={self.bias is not None}" + return s + + +class ColumnParallelLinear(LinearBase): + """Linear layer with column parallelism. + + The linear layer is defined as Y = XA + b. A is parallelized along + its second dimension as A = [A_1, ..., A_p]. + + Args: + input_size: first dimension of matrix A. + output_size: second dimension of matrix A. + bias: If true, add bias. + gather_output: If true, call all-gather on output and make Y available + to all GPUs, otherwise, every GPU will have its output + which is Y_i = XA_i + skip_bias_add: This was added to enable performance optimizations where + bias can be fused with other element-wise operations. we + skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + output_sizes: list of output sizes packed into one output, like for QKV + the list would be size 3. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = True, + gather_output: bool = False, + skip_bias_add: bool = False, + params_dtype: Optional[torch.dtype] = None, + quant_config: Optional[QuantizationConfig] = None, + output_sizes: Optional[List[int]] = None, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + use_presharded_weights: bool = False, + ): + super().__init__( + input_size, output_size, skip_bias_add, params_dtype, quant_config, prefix + ) + + self.gather_output = gather_output + self.use_presharded_weights = use_presharded_weights + + # Divide the weight matrix along the last dimension. + if tp_rank is None: + tp_rank = get_tensor_model_parallel_rank() + if tp_size is None: + tp_size = get_tensor_model_parallel_world_size() + self.tp_rank, self.tp_size = tp_rank, tp_size + assert self.quant_method is not None + self.output_size_per_partition = divide(self.output_size, tp_size) + self.output_partition_sizes = [self.output_size_per_partition] + # If QKV or MergedColumn, use output size of each partition. + if hasattr(self, "output_sizes"): + self.output_partition_sizes = [ + divide(output_size, tp_size) for output_size in self.output_sizes + ] + + if output_sizes is None: + output_sizes = [output_size] + + self.quant_method.create_weights( + layer=self, + input_size_per_partition=self.input_size, + output_partition_sizes=self.output_partition_sizes, + input_size=self.input_size, + output_size=self.output_size, + params_dtype=self.params_dtype, + weight_loader=( + self.weight_loader_v2 + if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED + else self.weight_loader + ), + ) + if bias: + self.bias = Parameter( + torch.empty(self.output_size_per_partition, dtype=params_dtype) + ) + set_weight_attrs( + self.bias, + { + "output_dim": 0, + "weight_loader": self.weight_loader, + }, + ) + else: + self.register_parameter("bias", None) + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + output_dim = getattr(param, "output_dim", None) + + # Special case for GGUF + is_gguf_weight = getattr(param, "is_gguf_weight", False) + is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False) + if is_gguf_weight_type: + param.weight_type = loaded_weight.item() + + # Materialize GGUF UninitializedParameter + if is_gguf_weight and isinstance(param, UninitializedParameter): + param.materialize(loaded_weight.shape, dtype=loaded_weight.dtype) + + use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False) + + param_data = param.data + # bitsandbytes loads the weights of the specific portion + # no need to narrow here + if output_dim is not None and not use_bitsandbytes_4bit: + shard_size = param_data.shape[output_dim] + start_idx = self.tp_rank * shard_size + + if _is_cpu: + from sglang.srt.model_loader.weight_utils import ( + narrow_padded_param_and_loaded_weight, + ) + + param_data, loaded_weight = narrow_padded_param_and_loaded_weight( + param_data, + loaded_weight, + 0, # param_data_start + start_idx, + output_dim, + shard_size, + not self.use_presharded_weights, + ) + else: + if not self.use_presharded_weights: + loaded_weight = loaded_weight.narrow( + output_dim, start_idx, shard_size + ) + + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def weight_loader_v2(self, param: Parameter, loaded_weight: torch.Tensor): + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + assert loaded_weight.numel() == 1 + loaded_weight = loaded_weight.reshape(1) + + if isinstance(param, _ColumnvLLMParameter): + param.load_column_parallel_weight( + loaded_weight, + tp_rank=self.tp_rank, + use_presharded_weights=self.use_presharded_weights, + ) + else: + # FIXME: This branch is needed to load deepseek v3 awq. + # However, we should fix this and avoid the branching here. + param.load_column_parallel_weight(loaded_weight) + + def forward(self, input_): + bias = self.bias if not self.skip_bias_add else None + + # Matrix multiply. + assert self.quant_method is not None + output_parallel = self.quant_method.apply(self, input_, bias) + if self.gather_output: + # All-gather across the partitions. + output = tensor_model_parallel_all_gather(output_parallel) + else: + output = output_parallel + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + def extra_repr(self) -> str: + s = f"in_features={self.input_size}" + s += f", output_features={self.output_size_per_partition}" + s += f", bias={self.bias is not None}" + s += f", tp_size={self.tp_size}" + s += f", gather_output={self.gather_output}" + return s + + +class MergedColumnParallelLinear(ColumnParallelLinear): + """Packed linear layers with column parallelism. + + Similar to ColumnParallelLinear, but the weight matrix is concatenated + along the output dimension. When the weight matrix is loaded, the + different partitions are sharded separately. + + Args: + input_size: input dimension of the linear layer. + output_sizes: list of output dimensions of the linear layer. + bias: If true, add bias. + gather_output: If true, call all-gather on output and make the output + available to all GPUs, otherwise, every GPU will have + its own output. + skip_bias_add: This was added to enable performance optimizations where + bias can be fused with other element-wise operations. we + skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__( + self, + input_size: int, + output_sizes: List[int], + bias: bool = True, + gather_output: bool = False, + skip_bias_add: bool = False, + params_dtype: Optional[torch.dtype] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + use_presharded_weights: bool = False, + ): + self.output_sizes = output_sizes + if tp_rank is None: + tp_rank = get_tensor_model_parallel_rank() + if tp_size is None: + tp_size = get_tensor_model_parallel_world_size() + self.tp_rank, self.tp_size = tp_rank, tp_size + assert all(output_size % tp_size == 0 for output_size in output_sizes) + self.use_presharded_weights = use_presharded_weights + super().__init__( + input_size=input_size, + output_size=sum(output_sizes), + bias=bias, + gather_output=gather_output, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + tp_rank=tp_rank, + tp_size=tp_size, + use_presharded_weights=use_presharded_weights, + ) + self.prefix = prefix + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: Optional[int] = None, + ): + + # Special case for GGUF + # initialize GGUF param after we know the quantize type + is_gguf_weight = getattr(param, "is_gguf_weight", False) + is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False) + if is_gguf_weight_type: + param.data[loaded_shard_id].copy_(loaded_weight) + param.shard_weight_type[loaded_shard_id] = loaded_weight.item() + return + + if is_gguf_weight: + output_dim = getattr(param, "output_dim", None) + shard_size = loaded_weight.size(output_dim) // self.tp_size + start_idx = self.tp_rank * shard_size + + loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) + + param.shard_id.append(loaded_shard_id) + param.shard_id_map[loaded_shard_id] = len(param.data_container) + param.data_container.append(loaded_weight) + return + + param_data = param.data + output_dim = getattr(param, "output_dim", None) + # Special case for AQLM codebooks. + is_metadata = getattr(param, "is_metadata", False) + # Special case for per-tensor scale to load scalar into fused array. + needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) + + if loaded_shard_id is None: + # Loaded weight is already fused on disk (qkv/mlp). + if output_dim is None: + if needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, 0 + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + return + current_shard_offset = 0 + shard_offsets: List[Tuple[int, int, int]] = [] + for i, output_size in enumerate(self.output_sizes): + shard_offsets.append((i, current_shard_offset, output_size)) + current_shard_offset += output_size + packed_dim = getattr(param, "packed_dim", None) + + use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False) + if _is_cpu: + shard_offsets = adjust_shard_offsets( + shard_offsets, loaded_weight, output_dim + ) + + for shard_id, shard_offset, shard_size in shard_offsets: + # Special case for Quantization. + # If quantized, we need to adjust the offset and size to account + # for the packing. + if packed_dim == output_dim: + shard_size = shard_size // param.pack_factor + shard_offset = shard_offset // param.pack_factor + # Special case for Marlin. + shard_size, shard_offset = adjust_marlin_shard( + param, shard_size, shard_offset + ) + + if use_bitsandbytes_4bit: + index = list(itertools.accumulate([0] + self.output_sizes)) + orig_offsets = { + str(i): (index[i], size) + for i, size in enumerate(self.output_sizes) + } + orig_offsets["total"] = (self.output_size, 0) + shard_size, shard_offset = adjust_bitsandbytes_4bit_shard( + param, orig_offsets, str(shard_id) + ) + + loaded_weight_shard = loaded_weight.narrow( + output_dim, shard_offset, shard_size + ) + self.weight_loader(param, loaded_weight_shard, shard_id) + return + + assert loaded_shard_id < len(self.output_sizes) + if output_dim is not None: + shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size + shard_size = self.output_sizes[loaded_shard_id] // self.tp_size + # Special case for quantization. + # If quantized, we need to adjust the offset and size to account + # for the packing. + packed_dim = getattr(param, "packed_dim", None) + if packed_dim == output_dim: + shard_size = shard_size // param.pack_factor + shard_offset = shard_offset // param.pack_factor + # Special case for Marlin. + shard_size, shard_offset = adjust_marlin_shard( + param, shard_size, shard_offset + ) + + use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False) + if use_bitsandbytes_4bit: + shard_size = loaded_weight.shape[output_dim] + shard_offset = loaded_weight.shape[output_dim] * loaded_shard_id + + param_data = param_data.narrow(output_dim, shard_offset, shard_size) + start_idx = self.tp_rank * shard_size + + if _is_cpu: + from sglang.srt.model_loader.weight_utils import ( + narrow_padded_param_and_loaded_weight, + ) + + param_data, loaded_weight = narrow_padded_param_and_loaded_weight( + param_data, + loaded_weight, + 0, # param_data_start + start_idx, + output_dim, + shard_size, + not use_bitsandbytes_4bit and not self.use_presharded_weights, + ) + else: + # bitsandbytes loads the weights of the specific portion + # no need to narrow here + if not use_bitsandbytes_4bit and not self.use_presharded_weights: + loaded_weight = loaded_weight.narrow( + output_dim, start_idx, shard_size + ) + + # Special case for AQLM codebooks. + elif is_metadata: + # metadata indicates fixed size concatenated along dim 0 + shard_size = loaded_weight.shape[0] + shard_offset = loaded_shard_id * shard_size + param_data = param_data.narrow(0, shard_offset, shard_size) + + # Special case for per-tensor scales in fused case. + elif needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, loaded_shard_id + ) + + else: + ignore_warning = getattr(param, "ignore_warning", False) + if not ignore_warning: + logger.warning( + "Loading a weight without `output_dim` attribute in " + "MergedColumnParallelLinear, assume the weight is " + "the same for all partitions." + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def _load_fused_module_from_checkpoint( + self, param: BasevLLMParameter, loaded_weight: torch.Tensor + ): + """ + Handle special case for models where MLP layers are already + fused on disk. In this case, we have no shard id. This function + determmines the shard id by splitting these layers and then calls + the weight loader using the shard id. + + An example of a model with these fused layers: + https://huggingface.co/microsoft/Phi-3-mini-4k-instruct + """ + + current_shard_offset = 0 + shard_offsets: List[Tuple[int, int, int]] = [] + for i, output_size in enumerate(self.output_sizes): + shard_offsets.append((i, current_shard_offset, output_size)) + current_shard_offset += output_size + + for shard_id, shard_offset, shard_size in shard_offsets: + # Special case for Quantization. + # If quantized, we need to adjust the offset and size to account + # for the packing. + if ( + isinstance(param, (PackedColumnParameter, PackedvLLMParameter)) + and param.packed_dim == param.output_dim + ): + shard_size, shard_offset = param.adjust_shard_indexes_for_packing( + shard_size=shard_size, shard_offset=shard_offset + ) + + loaded_weight_shard = loaded_weight.narrow( + param.output_dim, shard_offset, shard_size + ) + self.weight_loader_v2(param, loaded_weight_shard, shard_id) + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: Optional[int] = None, + ): + if loaded_shard_id is None: + if isinstance(param, PerTensorScaleParameter): + param.load_merged_column_weight( + loaded_weight=loaded_weight, + shard_id=0, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ) + return + elif type(param) in (RowvLLMParameter, BasevLLMParameter): + param.load_merged_column_weight( + loaded_weight=loaded_weight, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ) + return + # TODO: @dsikka - move to parameter.py + self._load_fused_module_from_checkpoint(param, loaded_weight) + return + + assert loaded_shard_id < len(self.output_sizes) + + if isinstance(param, BlockQuantScaleParameter): + weight_block_size = self.quant_method.quant_config.weight_block_size + block_n, _ = weight_block_size[0], weight_block_size[1] + shard_offset = ( + (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // block_n + ) // self.tp_size + shard_size = ( + (self.output_sizes[loaded_shard_id] + block_n - 1) + // block_n + // self.tp_size + ) + else: + shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size + shard_size = self.output_sizes[loaded_shard_id] // self.tp_size + + param.load_merged_column_weight( + loaded_weight=loaded_weight, + shard_id=loaded_shard_id, + shard_offset=shard_offset, + shard_size=shard_size, + use_presharded_weights=self.use_presharded_weights, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ) + + +class QKVParallelLinear(ColumnParallelLinear): + """Linear layers for the attention's QKV transformation. + + Linear layers for the linear transformation of the query, key, and value + vectors in the attention layer. The weight matrix is concatenated along + the output dimension. The layer is parallelized along the head dimension. + When the number of key/value heads is smaller than the number of query + heads (e.g., multi-query/grouped-query attention), the key/value head may + be replicated while the query heads are partitioned. + + Args: + hidden_size: input hidden state size of the transformer. + head_size: size of each attention head. + total_num_heads: total number of attention query heads. + total_num_kv_heads: total number of attention key/value heads. If + None, assume total_num_kv_heads = total_num_heads. + bias: If true, add bias. + skip_bias_add: This was added to enable performance optimizations where + bias can be fused with other element-wise operations. we + skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__( + self, + hidden_size: int, + head_size: int, + total_num_heads: int, + total_num_kv_heads: Optional[int] = None, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: Optional[torch.dtype] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + load_presharded_attn: bool = False, + ): + self.hidden_size = hidden_size + self.head_size = head_size + self.total_num_heads = total_num_heads + if total_num_kv_heads is None: + total_num_kv_heads = total_num_heads + self.total_num_kv_heads = total_num_kv_heads + # Divide the weight matrix along the last dimension. + if tp_rank is None: + tp_rank = get_tensor_model_parallel_rank() + if tp_size is None: + tp_size = get_tensor_model_parallel_world_size() + self.tp_rank, self.tp_size = tp_rank, tp_size + self.num_heads = divide(self.total_num_heads, tp_size) + if tp_size >= self.total_num_kv_heads: + self.num_kv_heads = 1 + self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads) + else: + self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) + self.num_kv_head_replicas = 1 + self.q_proj_shard_size = self.num_heads * self.head_size + self.kv_proj_shard_size = self.num_kv_heads * self.head_size + input_size = self.hidden_size + output_size = ( + (self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size + ) + self.output_sizes = [ + self.num_heads * self.head_size * tp_size, # q_proj + self.num_kv_heads * self.head_size * tp_size, # k_proj + self.num_kv_heads * self.head_size * tp_size, # v_proj + ] + self.use_presharded_weights = load_presharded_attn + + super().__init__( + input_size=input_size, + output_size=output_size, + bias=bias, + gather_output=False, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + tp_rank=tp_rank, + tp_size=tp_size, + use_presharded_weights=self.use_presharded_weights, + ) + + def _get_shard_offset_mapping(self, loaded_shard_id: str): + shard_offset_mapping = { + "q": 0, + "k": self.num_heads * self.head_size, + "v": (self.num_heads + self.num_kv_heads) * self.head_size, + "total": (self.num_heads + 2 * self.num_kv_heads) * self.head_size, + } + return shard_offset_mapping.get(loaded_shard_id) + + def _get_shard_size_mapping(self, loaded_shard_id: str): + shard_size_mapping = { + "q": self.num_heads * self.head_size, + "k": self.num_kv_heads * self.head_size, + "v": self.num_kv_heads * self.head_size, + } + return shard_size_mapping.get(loaded_shard_id) + + def _load_fused_module_from_checkpoint( + self, param: BasevLLMParameter, loaded_weight: torch.Tensor + ): + """ + Handle special case for models where QKV layers are already + fused on disk. In this case, we have no shard id. This function + determmines the shard id by splitting these layers and then calls + the weight loader using the shard id. + + An example of a model with these fused layers: + https://huggingface.co/microsoft/Phi-3-mini-4k-instruct + """ + shard_offsets = [ + # (shard_id, shard_offset, shard_size) + ("q", 0, self.total_num_heads * self.head_size), + ( + "k", + self.total_num_heads * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + ( + "v", + (self.total_num_heads + self.total_num_kv_heads) * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + ] + + for shard_id, shard_offset, shard_size in shard_offsets: + # Special case for Quantization. + # If quantized, we need to adjust the offset and size to account + # for the packing. + if ( + isinstance(param, (PackedColumnParameter, PackedvLLMParameter)) + and param.packed_dim == param.output_dim + ): + shard_size, shard_offset = param.adjust_shard_indexes_for_packing( + shard_size=shard_size, shard_offset=shard_offset + ) + + if not self.use_presharded_weights: + loaded_weight_shard = loaded_weight.narrow( + param.output_dim, shard_offset, shard_size + ) + self.weight_loader_v2(param, loaded_weight_shard, shard_id) + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: Optional[str] = None, + ): + if loaded_shard_id is None: # special case for certain models + if isinstance(param, PerTensorScaleParameter): + param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0) + return + elif type(param) in (RowvLLMParameter, BasevLLMParameter): + param.load_qkv_weight(loaded_weight=loaded_weight) + return + # TODO: @dsikka - move to parameter.py + self._load_fused_module_from_checkpoint(param, loaded_weight) + return + + assert loaded_shard_id in ["q", "k", "v"] + + shard_offset = self._get_shard_offset_mapping(loaded_shard_id) + shard_size = self._get_shard_size_mapping(loaded_shard_id) + + if isinstance(param, BlockQuantScaleParameter): + weight_block_size = self.quant_method.quant_config.weight_block_size + block_n, _ = weight_block_size[0], weight_block_size[1] + shard_offset = (shard_offset + block_n - 1) // block_n + shard_size = (shard_size + block_n - 1) // block_n + + param.load_qkv_weight( + loaded_weight=loaded_weight, + num_heads=self.num_kv_head_replicas, + shard_id=loaded_shard_id, + shard_offset=shard_offset, + shard_size=shard_size, + tp_rank=self.tp_rank, + use_presharded_weights=self.use_presharded_weights, + ) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: Optional[str] = None, + ): + + # Special case for GGUF + # initialize GGUF param after we know the quantize type + is_gguf_weight = getattr(param, "is_gguf_weight", False) + is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False) + if is_gguf_weight_type and loaded_shard_id is not None: + idx_map = {"q": 0, "k": 1, "v": 2} + param.data[idx_map[loaded_shard_id]].copy_(loaded_weight) + param.shard_weight_type[loaded_shard_id] = loaded_weight.item() + return + + if is_gguf_weight: + output_dim = getattr(param, "output_dim", None) + shard_size = loaded_weight.size(output_dim) // self.tp_size + start_idx = self.tp_rank * shard_size + + loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) + + param.shard_id.append(loaded_shard_id) + param.shard_id_map[loaded_shard_id] = len(param.data_container) + param.data_container.append(loaded_weight) + return + + param_data = param.data + output_dim = getattr(param, "output_dim", None) + # Special case for AQLM codebooks. + is_metadata = getattr(param, "is_metadata", False) + + # Special case for per-tensor scales in fused case. + needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) + + if loaded_shard_id is None: + # Loaded weight is already fused on disk (qkv/mlp). + if output_dim is None: + if needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, 0 + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + return + shard_offsets = [ + # (shard_id, shard_offset, shard_size) + ("q", 0, self.total_num_heads * self.head_size), + ( + "k", + self.total_num_heads * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + ( + "v", + (self.total_num_heads + self.total_num_kv_heads) * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + ] + use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False) + + packed_dim = getattr(param, "packed_dim", None) + if _is_cpu: + shard_offsets = adjust_shard_offsets( + shard_offsets, loaded_weight, output_dim + ) + + for shard_id, shard_offset, shard_size in shard_offsets: + # Special case for Quantized Weights. + # If quantized, we need to adjust the offset and size to account + # for the packing. + if packed_dim == output_dim: + shard_size = shard_size // param.pack_factor + shard_offset = shard_offset // param.pack_factor + + # Special case for Marlin. + shard_size, shard_offset = adjust_marlin_shard( + param, shard_size, shard_offset + ) + + if use_bitsandbytes_4bit: + orig_qkv_offsets = { + "q": (0, self.total_num_heads * self.head_size), + "k": ( + self.total_num_heads * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + "v": ( + (self.total_num_heads + self.total_num_kv_heads) + * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + "total": ( + (self.total_num_heads + 2 * self.total_num_kv_heads) + * self.head_size, + 0, + ), + } + + shard_size, shard_offset = adjust_bitsandbytes_4bit_shard( + param, orig_qkv_offsets, shard_id + ) + + if not self.use_presharded_weights: + loaded_weight_shard = loaded_weight.narrow( + output_dim, shard_offset, shard_size + ) + self.weight_loader(param, loaded_weight_shard, shard_id) + return + + assert loaded_shard_id in ["q", "k", "v"] + + # If output dim is defined, use the default loading process. + if output_dim is not None: + if loaded_shard_id == "q": + shard_offset = 0 + shard_size = self.num_heads * self.head_size + elif loaded_shard_id == "k": + shard_offset = self.num_heads * self.head_size + shard_size = self.num_kv_heads * self.head_size + elif loaded_shard_id == "v": + shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size + shard_size = self.num_kv_heads * self.head_size + # Special case for Quantized Weights. + # If quantized, we need to adjust the offset and size to account + # for the packing. + packed_dim = getattr(param, "packed_dim", None) + if packed_dim == output_dim: + shard_size = shard_size // param.pack_factor + shard_offset = shard_offset // param.pack_factor + + # Special case for Marlin. + shard_size, shard_offset = adjust_marlin_shard( + param, shard_size, shard_offset + ) + + use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False) + if use_bitsandbytes_4bit: + orig_qkv_offsets = { + "q": (0, self.num_heads * self.head_size), + "k": ( + self.num_heads * self.head_size, + self.num_kv_heads * self.head_size, + ), + "v": ( + (self.num_heads + self.num_kv_heads) * self.head_size, + self.num_kv_heads * self.head_size, + ), + "total": ( + (self.num_heads + 2 * self.num_kv_heads) * self.head_size, + 0, + ), + } + shard_size, shard_offset = adjust_bitsandbytes_4bit_shard( + param, orig_qkv_offsets, loaded_shard_id + ) + + param_data = param_data.narrow(output_dim, shard_offset, shard_size) + if loaded_shard_id == "q": + shard_id = self.tp_rank + else: + shard_id = self.tp_rank // self.num_kv_head_replicas + start_idx = shard_id * shard_size + + if _is_cpu: + from sglang.srt.model_loader.weight_utils import ( + narrow_padded_param_and_loaded_weight, + ) + + param_data, loaded_weight = narrow_padded_param_and_loaded_weight( + param_data, + loaded_weight, + 0, # param_data_start + start_idx, + output_dim, + shard_size, + not use_bitsandbytes_4bit and not self.use_presharded_weights, + ) + else: + # bitsandbytes loads the weights of the specific portion + # no need to narrow here + if not use_bitsandbytes_4bit and not self.use_presharded_weights: + loaded_weight = loaded_weight.narrow( + output_dim, start_idx, shard_size + ) + + # Special case for for AQLM codebooks. + elif is_metadata: + # metadata indicates fixed size concatenated along dim 0 + shard_size = loaded_weight.shape[0] + shard_index = ["q", "k", "v"].index(loaded_shard_id) + param_data = param_data.narrow(0, shard_index * shard_size, shard_size) + # Special case for per-tensor scales in fused case. + elif needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, loaded_shard_id + ) + else: + ignore_warning = getattr(param, "ignore_warning", False) + if not ignore_warning: + logger.warning( + "Loading a weight without `output_dim` attribute in " + "QKVParallelLinear, assume the weight is the same " + "for all partitions." + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class RowParallelLinear(LinearBase): + """Linear layer with row parallelism. + + The linear layer is defined as Y = XA + b. A is parallelized along + its first dimension and X along its second dimension as: + - - + | A_1 | + | . | + A = | . | X = [X_1, ..., X_p] + | . | + | A_p | + - - + Arguments: + input_size: first dimension of matrix A. + output_size: second dimension of matrix A. + bias: If true, add bias. Note that bias is not parallelized. + input_is_parallel: If true, we assume that the input is already + split across the GPUs and we do not split + again. + skip_bias_add: This was added to enable performance optimization where + bias can be fused with other element-wise operations. + We skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + """ + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = True, + input_is_parallel: bool = True, + skip_bias_add: bool = False, + params_dtype: Optional[torch.dtype] = None, + reduce_results: bool = True, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + use_presharded_weights: bool = False, + ): + super().__init__( + input_size, output_size, skip_bias_add, params_dtype, quant_config, prefix + ) + + self.input_is_parallel = input_is_parallel + self.reduce_results = reduce_results + + # Divide the weight matrix along the last dimension. + if tp_rank is None: + tp_rank = get_tensor_model_parallel_rank() + if tp_size is None: + tp_size = get_tensor_model_parallel_world_size() + self.tp_rank, self.tp_size = tp_rank, tp_size + self.input_size_per_partition = divide(input_size, self.tp_size) + assert self.quant_method is not None + self.use_presharded_weights = use_presharded_weights + + self.quant_method.create_weights( + layer=self, + input_size_per_partition=self.input_size_per_partition, + output_partition_sizes=[self.output_size], + input_size=self.input_size, + output_size=self.output_size, + params_dtype=self.params_dtype, + weight_loader=( + self.weight_loader_v2 + if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED + else self.weight_loader + ), + ) + + if bias: + self.bias = Parameter(torch.empty(self.output_size, dtype=params_dtype)) + set_weight_attrs( + self.bias, + { + "output_dim": 0, + "weight_loader": self.weight_loader, + }, + ) + else: + self.register_parameter("bias", None) + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + input_dim = getattr(param, "input_dim", None) + use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False) + + # Special case for GGUF + is_gguf_weight = getattr(param, "is_gguf_weight", False) + is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False) + if is_gguf_weight_type: + param.weight_type = loaded_weight.item() + + # Materialize GGUF UninitializedParameter + if is_gguf_weight and isinstance(param, UninitializedParameter): + weight_shape = list(loaded_weight.shape) + if input_dim: + weight_shape[input_dim] = weight_shape[input_dim] // self.tp_size + param.materialize(tuple(weight_shape), dtype=loaded_weight.dtype) + + param_data = param.data + # bitsandbytes loads the weights of the specific portion + # no need to narrow here + if ( + input_dim is not None + and not use_bitsandbytes_4bit + and not self.use_presharded_weights + ): + shard_size = param_data.shape[input_dim] + start_idx = self.tp_rank * shard_size + + if _is_cpu: + from sglang.srt.model_loader.weight_utils import ( + narrow_padded_param_and_loaded_weight, + ) + + param_data, loaded_weight = narrow_padded_param_and_loaded_weight( + param_data, + loaded_weight, + 0, # param_data_start + start_idx, + input_dim, + shard_size, + ) + else: + loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size) + + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor): + + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + assert loaded_weight.numel() == 1 + loaded_weight = loaded_weight.reshape(1) + + if isinstance(param, RowvLLMParameter): + # This `BasevLLMParameter` is defined in sglang/srt/layers/parameter.py, + # It supports additional parameters like tp_rank and use_presharded_weights. + param.load_row_parallel_weight( + loaded_weight, + tp_rank=self.tp_rank, + use_presharded_weights=self.use_presharded_weights, + ) + else: + # `params` is defined in `vllm/model_executor/parameter.py`, + # It does not support additional parameters. + param.load_row_parallel_weight(loaded_weight) + + def forward(self, input_, skip_all_reduce=False): + if self.input_is_parallel: + input_parallel = input_ + else: + splitted_input = split_tensor_along_last_dim( + input_, num_partitions=self.tp_size + ) + input_parallel = splitted_input[self.tp_rank].contiguous() + + # Matrix multiply. + assert self.quant_method is not None + # Only fuse bias add into GEMM for rank 0 (this ensures that + # bias will not get added more than once in TP>1 case) + bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias + with use_symmetric_memory(parallel_state.get_tp_group()) as sm: + output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_) + sm.tag(output_parallel) + + if self.reduce_results and self.tp_size > 1 and not skip_all_reduce: + output = tensor_model_parallel_all_reduce(output_parallel) + else: + output = output_parallel + + output_bias = self.bias if self.skip_bias_add else None + + return output, output_bias + + def extra_repr(self) -> str: + s = f"input_features={self.input_size_per_partition}" + s += f", output_features={self.output_size}" + s += f", bias={self.bias is not None}" + s += f", tp_size={self.tp_size}" + s += f", reduce_results={self.reduce_results}" + return s diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py new file mode 100644 index 000000000..f6603907a --- /dev/null +++ b/python/sglang/srt/layers/logits_processor.py @@ -0,0 +1,652 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Logits processing.""" + +import dataclasses +import logging +from typing import List, Optional, Union + +import torch +import triton +import triton.language as tl +from torch import nn + +from sglang.srt.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from sglang.srt.layers.dp_attention import ( + DpPaddingMode, + attn_tp_all_gather, + attn_tp_all_gather_into_tensor, + dp_gather_replicate, + dp_scatter, + get_attention_dp_rank, + get_attention_dp_size, + get_attention_tp_size, + get_global_dp_buffer, + get_local_attention_dp_size, + set_dp_buffer_len, +) +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardBatch, + ForwardMode, +) +from sglang.srt.utils import dump_to_file, is_npu, use_intel_amx_backend + +logger = logging.getLogger(__name__) + +_is_npu = is_npu() + + +@dataclasses.dataclass +class LogitsProcessorOutput: + ## Part 1: This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor + # The logits of the next tokens. shape: [#seq, vocab_size] + next_token_logits: torch.Tensor + # Used by speculative decoding (EAGLE) + # The last hidden layers + hidden_states: Optional[torch.Tensor] = None + + ## Part 2: This part will be assigned in python/sglang/srt/layers/sampler.py::Sampler + # he log probs of output tokens, if RETURN_ORIGINAL_LOGPROB = True, will get the log probs before applying temperature. If False, will get the log probs before applying temperature. + next_token_logprobs: Optional[torch.Tensor] = None + # The logprobs and ids of the top-k tokens in output positions. shape: [#seq, k] + next_token_top_logprobs_val: Optional[List] = None + next_token_top_logprobs_idx: Optional[List] = None + # The logprobs and ids of the requested token ids in output positions. shape: [#seq, n] (n is the number of requested token ids) + next_token_token_ids_logprobs_val: Optional[List] = None + next_token_token_ids_logprobs_idx: Optional[List] = None + + ## Part 3: Prefill-only. This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor + # The logprobs of input tokens. shape: [#token] + input_token_logprobs: Optional[torch.Tensor] = None + # The logprobs and ids of the top-k tokens in input positions. shape: [#seq, #token, k] + input_top_logprobs_val: List = None + input_top_logprobs_idx: List = None + # The logprobs and ids of the requested token ids in input positions. shape: [#seq, n] (n is the number of requested token ids) + input_token_ids_logprobs_val: Optional[List] = None + input_token_ids_logprobs_idx: Optional[List] = None + + +@dataclasses.dataclass +class LogitsMetadata: + forward_mode: ForwardMode + capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.NULL + next_token_logits_buffer: Optional[torch.Tensor] = None + + extend_return_logprob: bool = False + extend_return_top_logprob: bool = False + extend_token_ids_logprob: bool = False + extend_seq_lens: Optional[torch.Tensor] = None + extend_seq_lens_cpu: Optional[List[int]] = None + extend_logprob_start_lens_cpu: Optional[List[int]] = None + extend_logprob_pruned_lens_cpu: Optional[List[int]] = None + top_logprobs_nums: Optional[List[int]] = None + extend_input_logprob_token_ids_gpu: Optional[torch.Tensor] = None + token_ids_logprobs: Optional[List[List[int]]] = None + + # logits and logprobs post processing + temp_scaled_logprobs: bool = False + temperature: torch.Tensor = None + top_p_normalized_logprobs: bool = False + top_p: torch.Tensor = None + + # DP attention metadata. Not needed when DP attention is not used. + # Number of tokens in the request. + global_num_tokens_gpu: Optional[torch.Tensor] = None + # The start position of local hidden states. + dp_local_start_pos: Optional[torch.Tensor] = None + dp_local_num_tokens: Optional[torch.Tensor] = None + global_dp_buffer_len: Optional[int] = None + # Number of tokens to sample per DP rank + global_num_tokens_for_logprob_cpu: Optional[torch.Tensor] = None + global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None + # The gather mode for DP attention + dp_padding_mode: Optional[DpPaddingMode] = None + # for padding + padded_static_len: int = -1 + + @classmethod + def from_forward_batch(cls, forward_batch: ForwardBatch): + if ( + forward_batch.forward_mode.is_extend() + and forward_batch.return_logprob + and not forward_batch.forward_mode.is_target_verify() + ): + extend_return_top_logprob = any( + x > 0 for x in forward_batch.top_logprobs_nums + ) + extend_token_ids_logprob = any( + x is not None for x in forward_batch.token_ids_logprobs + ) + extend_return_logprob = False + extend_logprob_pruned_lens_cpu = [] + for extend_len, start_len in zip( + forward_batch.extend_seq_lens_cpu, + forward_batch.extend_logprob_start_lens_cpu, + ): + if extend_len - start_len > 0: + extend_return_logprob = True + extend_logprob_pruned_lens_cpu.append(extend_len - start_len) + else: + extend_return_logprob = extend_return_top_logprob = ( + extend_token_ids_logprob + ) = extend_logprob_pruned_lens_cpu = False + + return cls( + forward_mode=forward_batch.forward_mode, + capture_hidden_mode=forward_batch.capture_hidden_mode, + next_token_logits_buffer=forward_batch.next_token_logits_buffer, + extend_return_logprob=extend_return_logprob, + extend_return_top_logprob=extend_return_top_logprob, + extend_token_ids_logprob=extend_token_ids_logprob, + extend_seq_lens=forward_batch.extend_seq_lens, + extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu, + extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu, + extend_logprob_pruned_lens_cpu=extend_logprob_pruned_lens_cpu, + top_logprobs_nums=forward_batch.top_logprobs_nums, + token_ids_logprobs=forward_batch.token_ids_logprobs, + extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu, + padded_static_len=forward_batch.padded_static_len, + global_num_tokens_gpu=forward_batch.global_num_tokens_gpu, + dp_local_start_pos=forward_batch.dp_local_start_pos, + dp_local_num_tokens=forward_batch.dp_local_num_tokens, + global_dp_buffer_len=forward_batch.global_dp_buffer_len, + global_num_tokens_for_logprob_cpu=forward_batch.global_num_tokens_for_logprob_cpu, + global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu, + dp_padding_mode=DpPaddingMode.SUM_LEN, + ) + + def compute_dp_attention_metadata(self): + + cumtokens = torch.cumsum(self.global_num_tokens_for_logprob_gpu, dim=0) + dp_rank = get_attention_dp_rank() + if dp_rank == 0: + dp_local_start_pos = torch.zeros_like( + self.global_num_tokens_for_logprob_gpu[0] + ) + else: + dp_local_start_pos = cumtokens[dp_rank - 1] + dp_local_num_tokens = self.global_num_tokens_for_logprob_gpu[dp_rank] + + self.dp_local_start_pos = dp_local_start_pos + self.dp_local_num_tokens = dp_local_num_tokens + + if self.global_num_tokens_for_logprob_cpu is not None: + # create a smaller buffer to reduce peak memory usage + self.global_dp_buffer_len = sum(self.global_num_tokens_for_logprob_cpu) + else: + self.global_dp_buffer_len = self.global_dp_buffer_len + + set_dp_buffer_len( + self.global_dp_buffer_len, + self.dp_local_num_tokens, + self.global_num_tokens_for_logprob_cpu, + ) + + +class LogitsProcessor(nn.Module): + def __init__( + self, config, skip_all_gather: bool = False, logit_scale: Optional[float] = None + ): + super().__init__() + self.config = config + self.logit_scale = logit_scale + self.use_attn_tp_group = global_server_args_dict["enable_dp_lm_head"] + if self.use_attn_tp_group: + self.attn_tp_size = get_attention_tp_size() + self.do_tensor_parallel_all_gather = ( + not skip_all_gather and self.attn_tp_size > 1 + ) + self.do_tensor_parallel_all_gather_dp_attn = False + else: + self.do_tensor_parallel_all_gather = ( + not skip_all_gather and get_tensor_model_parallel_world_size() > 1 + ) + self.do_tensor_parallel_all_gather_dp_attn = ( + self.do_tensor_parallel_all_gather and get_attention_dp_size() != 1 + ) + self.final_logit_softcapping = getattr( + self.config, "final_logit_softcapping", None + ) + if ( + self.final_logit_softcapping is not None + and self.final_logit_softcapping < 0 + ): + self.final_logit_softcapping = None + + self.debug_tensor_dump_output_folder = global_server_args_dict.get( + "debug_tensor_dump_output_folder", None + ) + + def forward( + self, + input_ids, + hidden_states, + lm_head: VocabParallelEmbedding, + logits_metadata: Union[LogitsMetadata, ForwardBatch], + aux_hidden_states: Optional[torch.Tensor] = None, + ) -> LogitsProcessorOutput: + if isinstance(logits_metadata, ForwardBatch): + logits_metadata = LogitsMetadata.from_forward_batch(logits_metadata) + # Get the last hidden states and last logits for the next token prediction + if ( + logits_metadata.forward_mode.is_decode_or_idle() + or logits_metadata.forward_mode.is_target_verify() + ): + pruned_states = hidden_states + if aux_hidden_states is not None: + aux_pruned_states = [hidden for hidden in aux_hidden_states] + sample_indices = None + input_logprob_indices = None + elif ( + logits_metadata.forward_mode.is_extend() + and not logits_metadata.extend_return_logprob + ): + # Prefill without input logprobs. + if logits_metadata.padded_static_len < 0: + last_index = torch.cumsum(logits_metadata.extend_seq_lens, dim=0) - 1 + else: + # If padding_static length is 5 and extended_seq_lens is [2, 3], + # then our batch looks like [t00, t01, p, p, p, t10, t11, t12, p, p] + # and this retrieves t01 and t12, which are the valid last tokens + idx = torch.arange( + len(logits_metadata.extend_seq_lens), + device=logits_metadata.extend_seq_lens.device, + ) + last_index = ( + idx * logits_metadata.padded_static_len + + logits_metadata.extend_seq_lens + - 1 + ) + pruned_states = hidden_states[last_index] + if aux_hidden_states is not None: + aux_pruned_states = [hidden[last_index] for hidden in aux_hidden_states] + sample_indices = None + input_logprob_indices = None + else: + # Input logprobs are required. + # Find 3 different indices. + # 1. pruned_states: hidden states that we want logprobs from. + # 2. sample_indices: Indices that have sampled tokens. + # 3. input_logprob_indices: Indices that have input logprob tokens. + sample_index_pt = -1 + sample_indices = [] + input_logprob_indices_pt = 0 + input_logprob_indices = [] + pt, pruned_states = 0, [] + for extend_logprob_start_len, extend_len in zip( + logits_metadata.extend_logprob_start_lens_cpu, + logits_metadata.extend_seq_lens_cpu, + ): + # It can happen in chunked prefill. We still need to sample 1 token, + # But we don't want to include it in input logprob. + if extend_len == extend_logprob_start_len: + start_len = extend_logprob_start_len - 1 + else: + start_len = extend_logprob_start_len + + # We always need at least 1 token to sample because that's required + # by a caller. + assert extend_len > start_len + pruned_states.append(hidden_states[pt + start_len : pt + extend_len]) + pt += extend_len + sample_index_pt += extend_len - start_len + sample_indices.append(sample_index_pt) + input_logprob_indices.extend( + [ + input_logprob_indices_pt + i + for i in range(extend_len - extend_logprob_start_len) + ] + ) + input_logprob_indices_pt += extend_len - start_len + + pruned_states = torch.cat(pruned_states) + sample_indices = torch.tensor( + sample_indices, device=pruned_states.device, dtype=torch.int64 + ) + input_logprob_indices = torch.tensor( + input_logprob_indices, device=pruned_states.device, dtype=torch.int64 + ) + + # Compute logits for both input and sampled tokens. + logits = self._get_logits(pruned_states, lm_head, logits_metadata) + sampled_logits = ( + logits[sample_indices] if sample_indices is not None else logits + ) + + if self.debug_tensor_dump_output_folder: + assert ( + not self.do_tensor_parallel_all_gather + or get_local_attention_dp_size() == 1 + ), "dp attention + sharded lm_head doesn't support full logits" + full_logits = self._get_logits(hidden_states, lm_head, logits_metadata) + dump_to_file(self.debug_tensor_dump_output_folder, "logits", full_logits) + + hidden_states_to_store: Optional[torch.Tensor] = None + if logits_metadata.capture_hidden_mode.need_capture(): + if logits_metadata.capture_hidden_mode.is_full(): + if aux_hidden_states is not None: + aux_hidden_states = torch.cat(aux_hidden_states, dim=-1) + hidden_states_to_store = aux_hidden_states + else: + hidden_states_to_store = hidden_states + elif logits_metadata.capture_hidden_mode.is_last(): + # Get the last token hidden states. If sample_indices is None, + # pruned states only contain the last tokens already. + if aux_hidden_states is not None: + aux_pruned_states = torch.cat(aux_pruned_states, dim=-1) + hidden_states_to_store = ( + aux_pruned_states[sample_indices] + if sample_indices is not None + else aux_pruned_states + ) + else: + hidden_states_to_store = ( + pruned_states[sample_indices] + if sample_indices is not None + else pruned_states + ) + else: + assert False, "Should never reach" + + if not logits_metadata.extend_return_logprob: + # Decode mode or extend mode without return_logprob. + return LogitsProcessorOutput( + next_token_logits=sampled_logits, + hidden_states=hidden_states_to_store, + ) + else: + input_logprobs = logits[input_logprob_indices] + del hidden_states, logits + + # Normalize the logprob w/o temperature, top-p + pruned_lens = torch.tensor( + logits_metadata.extend_logprob_pruned_lens_cpu, + device=input_logprobs.device, + ) + if logits_metadata.temp_scaled_logprobs: + logits_metadata.temperature = torch.repeat_interleave( + logits_metadata.temperature.view(-1), + pruned_lens, + ).view(-1, 1) + if logits_metadata.top_p_normalized_logprobs: + logits_metadata.top_p = torch.repeat_interleave( + logits_metadata.top_p, + pruned_lens, + ) + input_logprobs = self.compute_temp_top_p_normalized_logprobs( + input_logprobs, logits_metadata + ) + + # Get the logprob of top-k tokens + if logits_metadata.extend_return_top_logprob: + ( + input_top_logprobs_val, + input_top_logprobs_idx, + ) = self.get_top_logprobs(input_logprobs, logits_metadata) + else: + input_top_logprobs_val = input_top_logprobs_idx = None + + # Get the logprob of given token id + if logits_metadata.extend_token_ids_logprob: + ( + input_token_ids_logprobs_val, + input_token_ids_logprobs_idx, + ) = self.get_token_ids_logprobs(input_logprobs, logits_metadata) + else: + input_token_ids_logprobs_val = input_token_ids_logprobs_idx = None + + input_token_logprobs = input_logprobs[ + torch.arange(input_logprobs.shape[0], device=input_logprobs.device), + logits_metadata.extend_input_logprob_token_ids_gpu, + ] + + return LogitsProcessorOutput( + next_token_logits=sampled_logits, + input_token_logprobs=input_token_logprobs, + input_top_logprobs_val=input_top_logprobs_val, + input_top_logprobs_idx=input_top_logprobs_idx, + hidden_states=hidden_states_to_store, + input_token_ids_logprobs_val=input_token_ids_logprobs_val, + input_token_ids_logprobs_idx=input_token_ids_logprobs_idx, + ) + + def _get_logits( + self, + hidden_states: torch.Tensor, + lm_head: VocabParallelEmbedding, + logits_metadata: LogitsMetadata, + embedding_bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Get logits from hidden_states. + + If sampled_logits_only is True, it means hidden_states only contain the + last position (e.g., extend without input logprobs). The caller should + guarantee the given hidden_states follow this constraint. + """ + if self.do_tensor_parallel_all_gather_dp_attn: + logits_metadata.compute_dp_attention_metadata() + hidden_states, local_hidden_states = ( + get_global_dp_buffer(), + hidden_states, + ) + dp_gather_replicate(hidden_states, local_hidden_states, logits_metadata) + + if hasattr(lm_head, "weight"): + if use_intel_amx_backend(lm_head): + logits = torch.ops.sgl_kernel.weight_packed_linear( + hidden_states.to(lm_head.weight.dtype), + lm_head.weight, + None, # bias + True, # is_vnni + ) + else: + logits = torch.matmul( + hidden_states.to(lm_head.weight.dtype), lm_head.weight.T + ) + else: + # GGUF models + # TODO: use weight_packed_linear for GGUF models + logits = lm_head.quant_method.apply(lm_head, hidden_states, embedding_bias) + + if self.logit_scale is not None: + logits.mul_(self.logit_scale) + + if self.do_tensor_parallel_all_gather: + if self.use_attn_tp_group: + if self.config.vocab_size % self.attn_tp_size == 0: + global_logits = torch.empty( + ( + self.attn_tp_size, + logits.shape[0], + self.config.vocab_size // self.attn_tp_size, + ), + device=logits.device, + dtype=logits.dtype, + ) + attn_tp_all_gather_into_tensor(global_logits, logits) + global_logits = global_logits.permute(1, 0, 2).reshape( + logits.shape[0], self.config.vocab_size + ) + else: + global_logits = torch.empty( + (self.config.vocab_size, logits.shape[0]), + device=logits.device, + dtype=logits.dtype, + ) + global_logits = global_logits.T + attn_tp_all_gather( + list(global_logits.tensor_split(self.attn_tp_size, dim=-1)), + logits, + ) + logits = global_logits + else: + logits = tensor_model_parallel_all_gather(logits) + + if self.do_tensor_parallel_all_gather_dp_attn: + logits, global_logits = ( + torch.empty( + (local_hidden_states.shape[0], logits.shape[1]), + device=logits.device, + dtype=logits.dtype, + ), + logits, + ) + dp_scatter(logits, global_logits, logits_metadata) + + if logits_metadata.next_token_logits_buffer is not None: + logits_buffer = logits_metadata.next_token_logits_buffer + assert logits_buffer.dtype == torch.float + logits_buffer.copy_(logits[:, : self.config.vocab_size]) + logits = logits_buffer + else: + logits = logits[:, : self.config.vocab_size].float() + + if self.final_logit_softcapping: + if not _is_npu: + fused_softcap(logits, self.final_logit_softcapping) + else: + logits = self.final_logit_softcapping * torch.tanh( + logits / self.final_logit_softcapping + ) + + return logits + + @staticmethod + def get_top_logprobs(all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata): + max_k = max(logits_metadata.top_logprobs_nums) + ret = all_logprobs.topk(max_k, dim=1) + values = ret.values.tolist() + indices = ret.indices.tolist() + + input_top_logprobs_val, input_top_logprobs_idx = [], [] + + pt = 0 + for k, pruned_len in zip( + logits_metadata.top_logprobs_nums, + logits_metadata.extend_logprob_pruned_lens_cpu, + ): + if pruned_len <= 0: + input_top_logprobs_val.append([]) + input_top_logprobs_idx.append([]) + continue + + input_top_logprobs_val.append( + [values[pt + j][:k] for j in range(pruned_len)] + ) + input_top_logprobs_idx.append( + [indices[pt + j][:k] for j in range(pruned_len)] + ) + pt += pruned_len + + return input_top_logprobs_val, input_top_logprobs_idx + + @staticmethod + def get_token_ids_logprobs( + all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata + ): + input_token_ids_logprobs_val, input_token_ids_logprobs_idx = [], [] + pt = 0 + for token_ids, pruned_len in zip( + logits_metadata.token_ids_logprobs, + logits_metadata.extend_logprob_pruned_lens_cpu, + ): + if pruned_len <= 0: + input_token_ids_logprobs_val.append([]) + input_token_ids_logprobs_idx.append([]) + continue + + input_token_ids_logprobs_val.append( + [all_logprobs[pt + j, token_ids].tolist() for j in range(pruned_len)] + ) + input_token_ids_logprobs_idx.append([token_ids for _ in range(pruned_len)]) + pt += pruned_len + + return input_token_ids_logprobs_val, input_token_ids_logprobs_idx + + @staticmethod + def compute_temp_top_p_normalized_logprobs( + last_logits: torch.Tensor, logits_metadata: LogitsMetadata + ) -> torch.Tensor: + """ + compute logprobs for the output token from the given logits. + + Returns: + torch.Tensor: logprobs from logits + """ + # Scale logits if temperature scaling is enabled + if logits_metadata.temp_scaled_logprobs: + last_logits = last_logits / logits_metadata.temperature + + # Normalize logprobs if top_p normalization is enabled + # NOTE: only normalize logprobs when top_p is set and not equal to 1.0 + if ( + logits_metadata.top_p_normalized_logprobs + and (logits_metadata.top_p != 1.0).any() + ): + from sglang.srt.layers.sampler import top_p_normalize_probs_torch + + probs = torch.softmax(last_logits, dim=-1) + del last_logits + probs = top_p_normalize_probs_torch(probs, logits_metadata.top_p) + return torch.log(probs) + else: + return torch.nn.functional.log_softmax(last_logits, dim=-1) + + +@triton.jit +def fused_softcap_kernel( + full_logits_ptr, + softcapping_value, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + # Load values + x = tl.load(full_logits_ptr + offsets, mask=mask) + + # Perform operations in-place + x = x / softcapping_value + + # Manual tanh implementation using exp + exp2x = tl.exp(2 * x) + x = (exp2x - 1) / (exp2x + 1) + + x = x * softcapping_value + + # Store result + tl.store(full_logits_ptr + offsets, x, mask=mask) + + +def fused_softcap(full_logits, final_logit_softcapping): + n_elements = full_logits.numel() + BLOCK_SIZE = 1024 + grid = ((n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE, 1, 1) + + fused_softcap_kernel[grid]( + full_logits_ptr=full_logits, + softcapping_value=final_logit_softcapping, + n_elements=n_elements, + BLOCK_SIZE=BLOCK_SIZE, + ) + return full_logits diff --git a/python/sglang/srt/layers/model_parallel.py b/python/sglang/srt/layers/model_parallel.py new file mode 100644 index 000000000..53c8b622e --- /dev/null +++ b/python/sglang/srt/layers/model_parallel.py @@ -0,0 +1,155 @@ +""" +Common utilities for torch model parallelism. +""" + +from typing import Optional, Sequence + +import torch +import torch.nn as nn +from torch.distributed.device_mesh import DeviceMesh + +try: + import torch.distributed.tensor as dt +except ImportError: + # torch 2.4 or older + import torch.distributed._tensor as dt + +from torch.distributed.tensor.parallel import ( + ColwiseParallel, + RowwiseParallel, + parallelize_module, +) + + +def _shard_tensor( + full_tensor: torch.Tensor, + device_mesh: DeviceMesh, + placements: Sequence[dt.Shard], +) -> "dt.DTensor": + """ + Locally shards a full tensor based on indicated sharding arrangement, and + returns a DTensor containing the local shard. + + .. warning:: This is a private API that is subject to change. It skips the + communication otherwise required by `distribute_tensor`. It is only + applicable to cases where all ranks have the same `full_tensor`. For + example, in distributed inference all ranks load from the same + checkpoint. This API will not check for data equality between ranks, it + is thus user's responsibility to ensure the `full_tensor` is the same + across ranks. + + Args: + full_tensor (torch.Tensor): the full tensor to be sharded. + device_mesh (:class:`DeviceMesh`): DeviceMesh to place the + DTensor. Must have same dimension as the number of placements. + placements (Sequence[:class:`Shard`]): the placements that + describes how to place the local tensor on DeviceMesh. + + Returns: + A :class:`DTensor` object with the shard as its local tensor. + + Examples: + >>> # xdoctest: +SKIP("need world_size and rank") + >>> device_mesh = dist.init_device_mesh("cuda", (world_size,)) + >>> full_tensor = torch.arange(world_size, device=f"cuda:{rank}") + >>> dtensor = _shard_tensor(full_tensor, device_mesh, [Shard(1)]) + """ + shape, offset = dt._utils.compute_local_shape_and_global_offset( + full_tensor.shape, device_mesh, placements + ) + slices = [ + slice(cur_offset, cur_offset + cur_shape) + for cur_shape, cur_offset in zip(shape, offset) + ] + local_tensor = full_tensor[slices] + return dt.DTensor.from_local(local_tensor, device_mesh, placements) + + +class ColwiseParallelSharded(ColwiseParallel): + """ + A version of ColwiseParallel where the local weight has been already + sharded. This is used for the fused wqkv case, where during loading, we + already sharded wq, wk, wv before fusing them. + """ + + # Override the _partition_linear_fn in ColwiseParallel + def _partition_linear_fn(self, name, module, device_mesh): + # colwise shard weight/bias to Shard(0), weight be Shard(0) + # means Colwise as Linear is input * weight^T + bias, where + # weight would become Shard(1) + for name, param in module.named_parameters(): + dtensor = dt.DTensor.from_local(param, device_mesh, [dt.Shard(0)]) + dist_param = torch.nn.Parameter(dtensor, requires_grad=False) + module.register_parameter(name, dist_param) + + +class RowwiseParallelMaybeWait(RowwiseParallel): + """ + A version of RowwiseParallel that waits for the output (establish dependency + between comm stream and compute stream in CUDA sense) before going into the + next op. This is needed to workaround the current interaction between + AsyncCollectiveTensor and custom ops, such as `class RMSNorm(CustomOp)`. + """ + + def _partition_linear_fn(self, name, module, device_mesh): + # Rowwise shard weight to Shard(1), bias to Replicate(), weight be Shard(1) + # means Rowwise as nn.Linear is input * weight^T + bias, where + # weight would become Shard(0) + module.register_parameter( + "weight", + nn.Parameter(_shard_tensor(module.weight, device_mesh, [dt.Shard(1)])), + ) + if getattr(module, "bias", None) is not None: + # The Linear module has bias + module.register_parameter( + "bias", + nn.Parameter( + dt.distribute_tensor(module.bias, device_mesh, [dt.Replicate()]) + ), + ) + + @staticmethod + def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): + outputs = super( + RowwiseParallelMaybeWait, RowwiseParallelMaybeWait + )._prepare_output_fn( + output_layouts, use_local_output, mod, outputs, device_mesh + ) + return torch.distributed._functional_collectives.wait_tensor(outputs) + + +def tensor_parallel( + module: torch.nn.Module, + device_mesh: Optional[DeviceMesh] = None, +): + """ + Tensor parallelize the model across the given device mesh. + Args: + module (`torch.nn.Module`): + The module to tensor parallelize. + device_mesh (`torch.distributed.DeviceMesh`): + The device mesh to use for tensor parallelism. + """ + + # Tensor parallelize a nn.Module based on the `_tp_plan` attribute of the module. + # No op if `_tp_plan` attribute does not exist under the module. + # This is a helper function to be used with `model.apply` to recursively + # parallelize a model. + def tplize(mod: torch.nn.Module) -> None: + tp_plan = getattr(mod, "_tp_plan", None) + if tp_plan is None: + return + for child_name, tp_style in tp_plan.items(): + submod = mod.get_submodule(child_name) + if tp_style == "Colwise": + parallelize_module(submod, device_mesh, ColwiseParallel()) + elif tp_style == "Rowwise": + parallelize_module(submod, device_mesh, RowwiseParallelMaybeWait()) + elif tp_style == "Colwise_Sharded": + parallelize_module(submod, device_mesh, ColwiseParallelSharded()) + else: + raise ValueError(f"Unknown TP style {tp_style}") + + # `apply` is a native method of `nn.Module` that recursively applies a + # function to every submodule. + module.apply(tplize) diff --git a/python/sglang/srt/layers/moe/__init__.py b/python/sglang/srt/layers/moe/__init__.py new file mode 100644 index 000000000..5c75a3682 --- /dev/null +++ b/python/sglang/srt/layers/moe/__init__.py @@ -0,0 +1,32 @@ +from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig +from sglang.srt.layers.moe.utils import ( + DeepEPMode, + MoeA2ABackend, + MoeRunnerBackend, + get_deepep_config, + get_deepep_mode, + get_moe_a2a_backend, + get_moe_runner_backend, + get_tbo_token_distribution_threshold, + initialize_moe_config, + is_tbo_enabled, + should_use_flashinfer_cutlass_moe_fp4_allgather, + should_use_flashinfer_trtllm_moe, +) + +__all__ = [ + "DeepEPMode", + "MoeA2ABackend", + "MoeRunner", + "MoeRunnerConfig", + "MoeRunnerBackend", + "initialize_moe_config", + "get_moe_a2a_backend", + "get_moe_runner_backend", + "get_deepep_mode", + "should_use_flashinfer_trtllm_moe", + "should_use_flashinfer_cutlass_moe_fp4_allgather", + "is_tbo_enabled", + "get_tbo_token_distribution_threshold", + "get_deepep_config", +] diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py new file mode 100755 index 000000000..d0fb4e3ef --- /dev/null +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -0,0 +1,365 @@ +"""CUTLASS based Fused MoE kernels.""" + +import torch + +from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() +if _is_cuda: + from sgl_kernel import ( + apply_shuffle_mul_sum, + cutlass_fp4_group_mm, + fp8_blockwise_scaled_grouped_mm, + prepare_moe_input, + scaled_fp4_experts_quant, + shuffle_rows, + silu_and_mul, + ) + + +def cutlass_fused_experts_fp8( + a: torch.Tensor, + w1_q: torch.Tensor, + w2_q: torch.Tensor, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + a1_strides: torch.Tensor, + c1_strides: torch.Tensor, + a2_strides: torch.Tensor, + c2_strides: torch.Tensor, + workspace: torch.Tensor, + a_ptrs: torch.Tensor, + b_ptrs: torch.Tensor, + out_ptrs: torch.Tensor, + a_scales_ptrs: torch.Tensor, + b_scales_ptrs: torch.Tensor, + expert_offsets: torch.Tensor, + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + use_fp8_blockscale: bool = True, +) -> torch.Tensor: + """Performs Fused MoE computation using CUTLASS-like kernels with FP8 weights and activations. + + This function implements a Mixture of Experts (MoE) layer with a SwiGLU/SiLU + activation, leveraging custom kernels likely derived from CUTLASS principles + for grouped matrix multiplication (`fp8_blockwise_scaled_grouped_mm`) and + data preparation (`prepare_moe_input`, `silu_and_mul`). + + It handles per-token routing, quantizes input activations to FP8 with + per-token scales, performs the expert computations using FP8 GEMMs with + pre-quantized FP8 weights (per-block scales), applies the SiLU activation, + and combines the results weighted by the router scores. + + Args: + a (torch.Tensor): Input activations. Shape: `(m, k)`, where `m` is the total + number of tokens and `k` is the hidden size. Expected dtype: `torch.half` + or `torch.bfloat16`. + w1_q (torch.Tensor): Pre-quantized FP8 weight tensor for the first GEMM + (up-projection part of SwiGLU). Expected shape: `(E, k, n*2)`, where + `E` is the number of experts, `k` is the hidden size, and `n*2` is the + intermediate size (`I`). Expected dtype: `torch.float8_e4m3fn`. + Note: This shape implies weights are stored as (num_experts, hidden_size, intermediate_size). + w2_q (torch.Tensor): Pre-quantized FP8 weight tensor for the second GEMM + (down-projection). Expected shape: `(E, n, k)`, where `n` is half the + intermediate size (`I // 2`). Expected dtype: `torch.float8_e4m3fn`. + Note: This shape implies weights are stored as (num_experts, intermediate_size // 2, hidden_size). + w1_scale (torch.Tensor): Scales corresponding to `w1_q` (per-block scales). + Shape: `(E, num_blocks_n, num_blocks_k)`. Dtype: `torch.float32`. + w2_scale (torch.Tensor): Scales corresponding to `w2_q` (per-block scales). + Shape: `(E, num_blocks_k, num_blocks_n)`. Dtype: `torch.float32`. + topk_weights (torch.Tensor): Router weights for the selected top-k experts + for each token. Shape: `(m, topk)`. Dtype should ideally match `a`. + topk_ids (torch.Tensor): Indices of the selected top-k experts for each token. + Shape: `(m, topk)`. Dtype: `torch.int32`. + a1_strides (torch.Tensor): Stride information for the first GEMM's 'a' input. + Passed directly to the underlying kernel. Expected shape `(E,)`, dtype `torch.int64`. + Note: Its exact usage within `fp8_blockwise_scaled_grouped_mm` needs clarification + as it's passed as both a_stride and b_stride in the first call. + c1_strides (torch.Tensor): Stride information for the first GEMM's 'c' output. + Passed directly to the underlying kernel. Expected shape `(E,)`, dtype `torch.int64`. + a2_strides (torch.Tensor): Stride information for the second GEMM's 'a' input. + Passed directly to the underlying kernel. Expected shape `(E,)`, dtype `torch.int64`. + Note: Its exact usage within `fp8_blockwise_scaled_grouped_mm` needs clarification + as it's passed as both a_stride and b_stride in the second call. + c2_strides (torch.Tensor): Stride information for the second GEMM's 'c' output. + Passed directly to the underlying kernel. Expected shape `(E,)`, dtype `torch.int64`. + workspace (torch.Tensor): Reusable workspace for the underlying kernel. + a_ptrs (torch.Tensor): Pointers container for calculating offsets of the input activations for each expert. + b_ptrs (torch.Tensor): Pointers container for calculating offsets of the input weights for each expert. + out_ptrs (torch.Tensor): Pointers container for calculating offsets of the output activations for each expert. + a_scales_ptrs (torch.Tensor): Pointers container for calculating offsets of the input scales for each expert. + b_scales_ptrs (torch.Tensor): Pointers container for calculating offsets of the input scales for each expert. + use_fp8_blockscale (bool, optional): Flag indicating usage of FP8 with + block scaling. Currently, only `True` is supported. Defaults to `True`. + + Returns: + torch.Tensor: The computed MoE layer output. Shape: `(m, k)`, dtype matches `a`. + + Raises: + AssertionError: If input shapes, dtypes, or flags are inconsistent or unsupported. + NotImplementedError: If CUDA is not available or `sgl_kernel` is not properly installed. + """ + assert use_fp8_blockscale, "Only support fp8 blockscale for now" + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" + assert w1_q.dtype == torch.float8_e4m3fn + assert w2_q.dtype == torch.float8_e4m3fn + assert a.shape[1] == w1_q.shape[1], "Hidden size mismatch w1" + assert w1_q.shape[2] == w2_q.shape[1] * 2, "Hidden size mismatch w2" + assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch" + assert w1_q.shape[0] == w2_q.shape[0], "Weights expert number mismatch" + assert w1_q.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch" + assert w1_q.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch" + assert a.dtype in [torch.half, torch.bfloat16], "Invalid output dtype" + + if is_cuda: + from sglang.srt.layers.quantization.fp8_kernel import ( + per_group_transpose, + per_token_group_quant_fp8_hopper_moe_mn_major, + sglang_per_token_group_quant_fp8, + ) + + out_dtype = a.dtype + num_experts = w1_q.size(0) + m = a.size(0) + k = w1_q.size(1) + n = w2_q.size(1) + + topk = topk_ids.size(1) + device = a.device + + a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) + c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) + + prepare_moe_input( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + a_map, + c_map, + num_experts, + n, + k, + ) + + a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128) + rep_a_q = shuffle_rows(a_q, a_map, (m * topk, k)) + rep_a1_scales = shuffle_rows(a1_scale, a_map, (m * topk, int(k / 128))) + + c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype) + c2 = torch.empty((m * topk, k), device=device, dtype=out_dtype) + + a_sf_layout = torch.empty((num_experts, 5), device=device, dtype=torch.int) + w_sf_layout = torch.empty((num_experts, 5), device=device, dtype=torch.int) + + fp8_blockwise_scaled_grouped_mm( + c1, + a_ptrs, + b_ptrs, + out_ptrs, + a_scales_ptrs, + b_scales_ptrs, + rep_a_q, + w1_q, + rep_a1_scales, + w1_scale, + a1_strides, + a1_strides, + c1_strides, + a_sf_layout, + w_sf_layout, + problem_sizes1, + expert_offsets[:-1], + workspace, + ) + + intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype) + silu_and_mul(c1, intermediate) + + intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128) + + fp8_blockwise_scaled_grouped_mm( + c2, + a_ptrs, + b_ptrs, + out_ptrs, + a_scales_ptrs, + b_scales_ptrs, + intemediate_q, + w2_q, + a2_scale, + w2_scale, + a2_strides, + a2_strides, + c2_strides, + a_sf_layout, + w_sf_layout, + problem_sizes2, + expert_offsets[:-1], + workspace, + ) + + result = torch.empty((m, k), device=device, dtype=out_dtype) + apply_shuffle_mul_sum(c2, result, c_map, topk_weights.to(out_dtype)) + return result + + +FLOAT4_E2M1_MAX = 6.0 +FLOAT8_E4M3_MAX = 448.0 + + +def cutlass_moe_fp4( + a: torch.Tensor, + a1_gscale: torch.Tensor, + w1_fp4: torch.Tensor, + w1_blockscale: torch.Tensor, + w1_alphas: torch.Tensor, + a2_gscale: torch.Tensor, + w2_fp4: torch.Tensor, + w2_blockscale: torch.Tensor, + w2_alphas: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + params: CutlassMoEParams, + apply_router_weight_on_input: bool = False, +): + """ + MoE implementation for FP4 Inputs + + # Gemm 1 + a: Input tensor: [m, k] (half/bfloat16) + a1_gscale: Activation scale per expert: [e] (float32) + w1(gate up) (not an argument to cutlass_moe_fp4): [e, 2 * n, k] + w1_fp4: [e, 2 * n, k // 2], dtype: torch.uint8 (stacked fp4: E2M1) + (Note: `n` is the up projection output dim, `k` is the input dim in + full precision) + w1_blockscale: [e, 2 * n, k // block_size] (float8_e4m3) + (Block size = 16 for NVFP4) + + # Gemm 2 + a2_gscale: Activation scale per expert: [e] + w2(down projection) (not an argument to cutlass_moe_fp4): [e, k, n] + w2_fp4: [e, k, n // 2], dtype: torch.uint8 (stacked E2M1) + w2_blockscale: [e, k, n // block_size], dtype: float8_e4m3 + + Strides for activations, weights and output in logical number of elements. + The activations & output stride is the number of elements to the next row. + The weights stride is the number of elements to the next row per expert. + For example, if the weight is [e, n, k], then the b_stride is a tensor of + shape [e] with each element being k. Similarly for activations, if the + shape is [m, k], then the a_stride has shape [e] with each value k. + Similarly for output, if the output is [m, n], then the c_stride is a + tensor of shape [e] with each element being k. + + Note: cutlass_fp4_group_mm is designed to accept the strides of + activations and weights to be the same, so it is passed in as a single + tensor. + ab_strides_13: [e] dtype: int64 [Gemm 1: Activation / Weight strides] + ab_strides_2: [e] dtype: int64 [Gemm 2: Activation / Weight strides] + c_strides_13: [e] dtype: int64 [Gemm 1: Output Strides] + c_strides_2: [e] dtype: int64 [Gemm 1: Output Strides] + + topk_weights: [m, topk] dtype: float8 + topk_ids: [m, topk] dtype: float8 + + m, n, k: Unquantized weight shapes, dtype: int + e: number of experts for the current rank, dtype: int + assumes that topk < k < n to satisfy - up/down projection expectations. + """ + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" + assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8" + assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8" + assert ( + w1_fp4.ndim == 3 + and w2_fp4.ndim == 3 + and w1_blockscale.ndim == 3 + and w2_blockscale.ndim == 3 + ), "All Weights must be of rank 3 for cutlass_moe_fp4" + m_a, k_a = a.shape + e_w1, nx2_w1, half_k_w1 = w1_fp4.shape + e_w2, k_w2, half_n_w2 = w2_fp4.shape + + assert e_w1 == e_w2 and e_w1 == params.num_experts, ( + "Number of experts must match", + " between weights.", + ) + assert ( + k_a // 2 == half_k_w1 and params.hidden_size == k_w2 + ), "Hidden size mismatch between a, w1 and w2" + assert ( + nx2_w1 == params.intermediate_size_per_partition * 2 + and half_n_w2 == params.intermediate_size_per_partition // 2 + ), ("mismatch in " "expected `n`") + assert 2 * half_k_w1 == k_w2, "Hidden size mismatch w2 and w1" + assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype" + + out_dtype = a.dtype + num_topk = topk_ids.shape[1] + device = a.device + a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) + c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) + prepare_moe_input( + topk_ids, + params.expert_offsets, + params.problem_sizes1, + params.problem_sizes2, + a_map, + c_map, + params.num_experts, + params.intermediate_size_per_partition, + params.hidden_size, + params.blockscale_offsets, + ) + + rep_a_fp4, rep_a_blockscale = scaled_fp4_experts_quant( + a, + a1_gscale, + params.expert_offsets, + params.blockscale_offsets, + num_topk, + expert_map=a_map, + ) + c1 = cutlass_fp4_group_mm( + rep_a_fp4, + w1_fp4, + rep_a_blockscale, + w1_blockscale, + w1_alphas, + out_dtype, + device, + params.to_gemm1_args(), + ) + del rep_a_fp4, rep_a_blockscale + + # hidden size dimension is split to one halfpytho sized tensor. + intermediate = torch.empty( + (m_a * num_topk, w1_fp4.shape[1] // 2), device=device, dtype=out_dtype + ) + silu_and_mul(c1, intermediate) + + int_fp4, int_blockscale = scaled_fp4_experts_quant( + intermediate, + a2_gscale, + params.expert_offsets, + params.blockscale_offsets, + num_topk, + ) + c2 = cutlass_fp4_group_mm( + int_fp4, + w2_fp4, + int_blockscale, + w2_blockscale, + w2_alphas, + out_dtype, + device, + params.to_gemm2_args(), + ) + del int_fp4, int_blockscale + c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size)) + c2 = c2.view(m_a, num_topk, params.hidden_size) + if not apply_router_weight_on_input: + c2 = c2 * topk_weights.view(m_a, num_topk, 1).to(out_dtype) + return c2.sum(dim=1).to(out_dtype) diff --git a/python/sglang/srt/layers/moe/cutlass_moe_params.py b/python/sglang/srt/layers/moe/cutlass_moe_params.py new file mode 100644 index 000000000..f3de60e04 --- /dev/null +++ b/python/sglang/srt/layers/moe/cutlass_moe_params.py @@ -0,0 +1,169 @@ +from dataclasses import dataclass +from enum import Enum, auto +from typing import Optional + +import torch + + +class CutlassMoEType(Enum): + """ + Enum for the different types of cutlass moe operations + that are currently supported in SGLang. + """ + + BlockscaledFP8 = auto() + BlockscaledFP4 = auto() + + +@dataclass +class CutlassMoEParams: + """ + Parameters for the cutlass moe operation. + """ + + # Type as defined above + cutlass_moe_type: CutlassMoEType + + # Strides for activations, weights and output in logical number of elements. + # The activations & output stride is the number of elements to the next row. + # The weights stride is the number of elements to the next row per expert. + # For example, if the weight is [e, n, k], then the b_stride is a tensor of + # shape [e] with each element being k. Similarly for activations, if the + # shape is [m, k], then the a_stride has shape [e] with each value k. + # Similarly for output, if the output is [m, n], then the c_stride is a + # tensor of shape [e] with each element being k. + + # Note: cutlass_fp4_group_mm is designed to accept the strides of + # activations and weights to be the same, so it is passed in as a single + # tensor. + # ab_strides_13: [e] dtype: int64 [Gemm 1: Activation / Weight strides] + # ab_strides_2: [e] dtype: int64 [Gemm 2: Activation / Weight strides] + # c_strides_13: [e] dtype: int64 [Gemm 1: Output Strides] + # c_strides_2: [e] dtype: int64 [Gemm 2: Output Strides] + ab_strides_13: torch.Tensor + ab_strides_2: torch.Tensor + c_strides_13: torch.Tensor + c_strides_2: torch.Tensor + + # m: Total number of tokens + # n: intermediate size per partition + # k: hidden size per expert + # e: Number of experts + # device: Device to run computation on and store tensors + m: int + intermediate_size_per_partition: int + hidden_size: int + num_experts: int + device: torch.device + + # Pointers container for calculating offsets of the input activations for each expert + # a_ptrs: [e] dtype: int64 + a_ptrs: torch.Tensor + + # Pointers container for calculating offsets of the input weights for each expert + # b_ptrs: [e] dtype: int64 + b_ptrs: torch.Tensor + + # Pointers container for calculating offsets of the output activations for each expert + # out_ptrs: [e] dtype: int64 + out_ptrs: torch.Tensor + # Pointers container for calculating offsets of the input scales for each expert + # a_scales_ptrs: [e] dtype: int64 + # b_scales_ptrs: [e] dtype: int64 + a_scales_ptrs: torch.Tensor + b_scales_ptrs: torch.Tensor + + # Offsets that mark at which token index each expert begins its computation + # The number of tokens computed with expert E is expert_offsets[E + 1] - expert_offsets[E] + # expert_offsets: [e+1] dtype: int32 + expert_offsets: torch.Tensor + + # Problem size: (num_experts, (m,2n,k)) for first GEMM + # problem_sizes1: [e, 3] dtype: int32 + # Problem size: (num_experts, (m,n,k)) for second GEMM + # problem_sizes2: [e, 3] dtype: int32 + problem_sizes1: torch.Tensor + problem_sizes2: torch.Tensor + # Similar to expert_offsets, but for blockscales for FP4 blockscaled Group GEMM + blockscale_offsets: Optional[torch.Tensor] = None + + def __init__( + self, + cutlass_moe_type: CutlassMoEType, + device: torch.device, + num_experts: int, + intermediate_size_per_partition: int, + hidden_size: int, + ): + self.cutlass_moe_type = cutlass_moe_type + self.device = device + self.num_experts = num_experts + self.intermediate_size_per_partition = intermediate_size_per_partition + self.hidden_size = hidden_size + self.n = self.intermediate_size_per_partition + self.k = self.hidden_size + self.e = self.num_experts + self.ab_strides_13 = torch.full( + (self.e,), self.k, dtype=torch.int64, device=self.device + ) + self.ab_strides_2 = torch.full( + (self.e,), self.n, dtype=torch.int64, device=self.device + ) + self.c_strides_13 = torch.full( + (self.e,), 2 * self.n, dtype=torch.int64, device=self.device + ) + self.c_strides_2 = torch.full( + (self.e,), self.k, dtype=torch.int64, device=self.device + ) + self.expert_offsets = torch.empty( + (self.e + 1,), dtype=torch.int32, device=self.device + ) + self.problem_sizes1 = torch.empty( + (self.e, 3), dtype=torch.int32, device=self.device + ) + self.problem_sizes2 = torch.empty( + (self.e, 3), dtype=torch.int32, device=self.device + ) + if self.cutlass_moe_type == CutlassMoEType.BlockscaledFP4: + self.blockscale_offsets = torch.empty( + (self.e + 1,), dtype=torch.int32, device=self.device + ) + else: + self.blockscale_offsets = None + self.a_ptrs = torch.empty((self.e,), dtype=torch.int64, device=self.device) + self.b_ptrs = torch.empty((self.e,), dtype=torch.int64, device=self.device) + self.out_ptrs = torch.empty((self.e,), dtype=torch.int64, device=self.device) + self.a_scales_ptrs = torch.empty( + (self.e,), dtype=torch.int64, device=self.device + ) + self.b_scales_ptrs = torch.empty( + (self.e,), dtype=torch.int64, device=self.device + ) + + def to_gemm1_args(self) -> dict: + return { + "ab_strides": self.ab_strides_13, + "c_strides": self.c_strides_13, + "problem_sizes": self.problem_sizes1, + "expert_offsets": self.expert_offsets[:-1], + "blockscale_offsets": self.blockscale_offsets[:-1], + # "a_ptrs": self.a_ptrs, + # "b_ptrs": self.b_ptrs, + # "out_ptrs": self.out_ptrs, + # "a_scales_ptrs": self.a_scales_ptrs, + # "b_scales_ptrs": self.b_scales_ptrs, + } + + def to_gemm2_args(self) -> dict: + return { + "ab_strides": self.ab_strides_2, + "c_strides": self.c_strides_2, + "problem_sizes": self.problem_sizes2, + "expert_offsets": self.expert_offsets[:-1], + "blockscale_offsets": self.blockscale_offsets[:-1], + # "a_ptrs": self.a_ptrs, + # "b_ptrs": self.b_ptrs, + # "out_ptrs": self.out_ptrs, + # "a_scales_ptrs": self.a_scales_ptrs, + # "b_scales_ptrs": self.b_scales_ptrs, + } diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py new file mode 100644 index 000000000..216424eea --- /dev/null +++ b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cutlass W4A8 MoE kernel.""" +from typing import Optional + +import torch +from sgl_kernel import ( + cutlass_w4a8_moe_mm, + get_cutlass_w4a8_moe_mm_data, + sgl_per_tensor_quant_fp8, + silu_and_mul, +) + +from sglang.srt.layers.moe.ep_moe.kernels import ( + post_reorder_triton_kernel_for_cutlass_moe, + pre_reorder_triton_kernel_for_cutlass_moe, + run_cutlass_moe_ep_preproess, +) + + +def cutlass_w4a8_moe( + start_expert_id: int, + end_expert_id: int, + total_num_experts: int, + a: torch.Tensor, + w1_q: torch.Tensor, + w2_q: torch.Tensor, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids_: torch.Tensor, + local_topk_ids: torch.Tensor, + a_strides1: torch.Tensor, + b_strides1: torch.Tensor, + c_strides1: torch.Tensor, + a_strides2: torch.Tensor, + b_strides2: torch.Tensor, + c_strides2: torch.Tensor, + s_strides13: torch.Tensor, + s_strides2: torch.Tensor, + expert_offsets: torch.Tensor, + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + apply_router_weight_on_input: bool = False, +) -> torch.Tensor: + """ + This function computes a w4a8-quantized Mixture of Experts (MoE) layer + using two sets of quantized weights, w1_q and w2_q, and top-k gating + mechanism. The matrix multiplications are implemented with CUTLASS + grouped gemm. + + Parameters: + - a (torch.Tensor): The input tensor to the MoE layer. + Shape: [M, K] + - w1_q (torch.Tensor): The first set of int4-quantized expert weights. + Shape: [num_experts, N * 2, K // 2] + (the weights are passed transposed and int4-packed) + - w2_q (torch.Tensor): The second set of int4-quantized expert weights. + Shape: [num_experts, K, N // 2] + (the weights are passed transposed and int4-packed) + - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q. + Shape: [num_experts, K // 512, N * 8] + - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q. + Shape: [num_experts, N // 512, K * 4] + - topk_weights (torch.Tensor): The weights of each token->expert mapping. + - a_strides1 (torch.Tensor): The input strides of the first grouped gemm. + - b_strides1 (torch.Tensor): The weights strides of the first grouped gemm. + - c_strides1 (torch.Tensor): The output strides of the first grouped gemm. + - a_strides2 (torch.Tensor): The input strides of the second grouped gemm. + - b_strides2 (torch.Tensor): The weights strides of the second grouped gemm. + - c_strides2 (torch.Tensor): The output strides of the second grouped gemm. + - s_strides13 (torch.Tensor): The input and scale strides of the first grouped gemm. + - s_strides2 (torch.Tensor): The scale strides of the second grouped gemm. + - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a. + Shape: scalar or [1, K] + - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to + quantize the intermediate result between the gemms. + Shape: scalar or [1, N] + - apply_router_weight_on_input (bool): When true, the topk weights are + applied directly on the inputs. This is only applicable when topk is 1. + + Returns: + - torch.Tensor: The fp8 output tensor after applying the MoE layer. + """ + assert topk_weights.shape == topk_ids_.shape, "topk shape mismatch" + assert w1_q.dtype == torch.int8 + assert w2_q.dtype == torch.int8 + assert a.shape[1] // 2 == w1_q.shape[2], "Hidden size mismatch w1" + assert w1_q.shape[2] * 2 == w2_q.shape[1], "Hidden size mismatch w2" + assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch" + assert w1_q.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch" + assert w1_q.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch" + + assert a_strides1.shape[0] == w1_q.shape[0], "A Strides 1 expert number mismatch" + assert b_strides1.shape[0] == w1_q.shape[0], "B Strides 1 expert number mismatch" + assert a_strides2.shape[0] == w2_q.shape[0], "A Strides 2 expert number mismatch" + assert b_strides2.shape[0] == w2_q.shape[0], "B Strides 2 expert number mismatch" + num_experts = w1_q.size(0) + m = a.size(0) + k = w1_q.size(2) * 2 # w1_q is transposed and packed + n = w2_q.size(2) * 2 # w2_q is transposed and packed + topk = topk_ids_.size(1) + + if apply_router_weight_on_input: + assert topk == 1, "apply_router_weight_on_input is only implemented for topk=1" + + device = a.device + + _, src2dst, _ = run_cutlass_moe_ep_preproess( + local_topk_ids, + num_experts, + ) + + gateup_input = torch.empty( + (m * topk, k), + device=device, + dtype=torch.float8_e4m3fn, + ) + + pre_reorder_triton_kernel_for_cutlass_moe[(m,)]( + a, + gateup_input, + src2dst, + local_topk_ids, + a1_scale, + total_num_experts, + topk, + k, + BLOCK_SIZE=512, + ) + + # NOTE: a_map and c_map are not used in the get_cutlass_w4a8_moe_mm_data kernel, + # they are kept to allow for a quick switch of the permutation logic + # from the current triton kernel implementation to the cutlass-based one if needed. + a_map = torch.empty((local_topk_ids.numel()), dtype=torch.int32, device=device) + c_map = torch.empty((local_topk_ids.numel()), dtype=torch.int32, device=device) + get_cutlass_w4a8_moe_mm_data( + local_topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + a_map, + c_map, + num_experts, + n, + k, + ) + + c1 = torch.empty((m * topk, n * 2), device=device, dtype=torch.bfloat16) + c2 = torch.zeros((m * topk, k), device=device, dtype=torch.bfloat16) + + cutlass_w4a8_moe_mm( + c1, + gateup_input, + w1_q, + a1_scale.float(), + w1_scale, + expert_offsets[:-1], + problem_sizes1, + a_strides1, + b_strides1, + c_strides1, + s_strides13, + 128, + topk, + ) + + intermediate = torch.empty((m * topk, n), device=device, dtype=torch.bfloat16) + silu_and_mul(c1, intermediate) + + intermediate_q = torch.empty( + intermediate.shape, dtype=torch.float8_e4m3fn, device=device + ) + sgl_per_tensor_quant_fp8(intermediate, intermediate_q, a2_scale.float(), True) + + cutlass_w4a8_moe_mm( + c2, + intermediate_q, + w2_q, + a2_scale.float(), + w2_scale, + expert_offsets[:-1], + problem_sizes2, + a_strides2, + b_strides2, + c_strides2, + s_strides2, + 128, + topk, + ) + + output = torch.empty_like(a) + post_reorder_triton_kernel_for_cutlass_moe[(m,)]( + c2, + output, + src2dst, + local_topk_ids, + topk_weights, + num_experts, + topk, + k, + 0, + BLOCK_SIZE=512, + ) + return output diff --git a/python/sglang/srt/layers/moe/ep_moe/__init__.py b/python/sglang/srt/layers/moe/ep_moe/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/layers/moe/ep_moe/kernels.py b/python/sglang/srt/layers/moe/ep_moe/kernels.py new file mode 100644 index 000000000..08660812d --- /dev/null +++ b/python/sglang/srt/layers/moe/ep_moe/kernels.py @@ -0,0 +1,1438 @@ +import logging +from typing import List, Optional + +import torch +import triton + +from sglang.srt.layers.quantization.fp8_kernel import per_token_group_quant_fp8 +from sglang.srt.utils import ceil_div, dispose_tensor, is_cuda +from sglang.utils import is_in_ci + +logger = logging.getLogger(__name__) + +_is_cuda = is_cuda() +if _is_cuda: + from sglang.srt.layers.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8 as per_token_group_quant_fp8, + ) + +import triton.language as tl + + +@triton.jit +def deepep_permute_triton_kernel( + input_ptr, + gateup_input_ptr, + src2dst_ptr, + topk_ids_ptr, + a1_scales_ptr, + topk, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + OutDtype = gateup_input_ptr.dtype.element_ty + + src_idx = tl.program_id(0) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + + src_ptr = input_ptr + src_idx * hidden_size + + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + in_data = tl.load(src_ptr + offset, mask=mask).to(OutDtype) + + for idx in range(topk): + dst_idx = tl.load(src2dst_ptr + idx) + if dst_idx >= 0: + dst_ptr = gateup_input_ptr + dst_idx * hidden_size + tl.store(dst_ptr + offset, in_data, mask=mask) + + +@triton.jit +def deepep_post_reorder_triton_kernel( + down_output_ptr, + output_ptr, + src2dst_ptr, + topk_ids_ptr, + topk_weights_ptr, + topk, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + InDtype = down_output_ptr.dtype.element_ty + + src_idx = tl.program_id(0) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + topk_weights_ptr = topk_weights_ptr + src_idx * topk + + store_ptr = output_ptr + src_idx * hidden_size + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype) + for idx in range(topk): + dst_idx = tl.load(src2dst_ptr + idx) + if dst_idx >= 0: + weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype) + load_ptr = down_output_ptr + dst_idx * hidden_size + in_data = tl.load(load_ptr + offset, mask=mask) + sum_vec += in_data * weigh_scale + tl.store(store_ptr + offset, sum_vec, mask=mask) + + +@triton.jit +def compute_src2dst_triton_kernel( + reorder_ids, src2dst, num_toks, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(axis=0) + dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = dst_id < num_toks + src_id = tl.load(reorder_ids + dst_id, mask=mask) + tl.store(src2dst + src_id, dst_id, mask=mask) + + +@triton.jit +def deepep_compute_src2dst_triton_kernel( + reorder_ids, src2dst, num_toks, num_minus_one, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(axis=0) + dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = dst_id < num_toks + src_id = tl.load(reorder_ids + dst_id, mask=mask) + num_invalid = tl.load(num_minus_one) + tl.store(src2dst + src_id, dst_id - num_invalid, mask=mask) + + +def deepep_run_moe_deep_preprocess(topk_ids: torch.Tensor, num_experts: int): + reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True) + seg_indptr = torch.empty(num_experts + 1, device=topk_ids.device, dtype=torch.int64) + src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int64) + + # Find offset + expert_ids = torch.arange( + num_experts + 1, device=topk_ids.device, dtype=reorder_topk_ids.dtype + ) + torch.searchsorted(reorder_topk_ids, expert_ids, out=seg_indptr) + num_minus_one = seg_indptr[0] + seg_indptr = seg_indptr - num_minus_one + + BLOCK_SIZE = 512 + grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),) + deepep_compute_src2dst_triton_kernel[grid]( + reorder_ids, src2dst, topk_ids.numel(), num_minus_one, BLOCK_SIZE + ) + reorder_topk_ids = reorder_topk_ids[num_minus_one:] + return reorder_topk_ids, src2dst, seg_indptr + + +@triton.jit +def compute_seg_indptr_triton_kernel(reorder_topk_ids, seg_indptr, num_toks): + expert = tl.program_id(0) + low = 0 + high = num_toks - 1 + target_location = -1 + while low <= high: + mid = (low + high) // 2 + + if tl.load(reorder_topk_ids + mid) > expert: + high = mid - 1 + else: + low = mid + 1 + target_location = mid + tl.store(seg_indptr + expert + 1, target_location + 1) + + +def run_moe_ep_preproess(topk_ids: torch.Tensor, num_experts: int): + reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True) + + seg_indptr = torch.zeros(num_experts + 1, device=topk_ids.device, dtype=torch.int64) + src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int32) + + compute_seg_indptr_triton_kernel[(num_experts,)]( + reorder_topk_ids, seg_indptr, topk_ids.numel() + ) + + BLOCK_SIZE = 512 + grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),) + compute_src2dst_triton_kernel[grid]( + reorder_ids, src2dst, topk_ids.numel(), BLOCK_SIZE + ) + + return reorder_topk_ids, src2dst, seg_indptr + + +def run_cutlass_moe_ep_preproess(local_topk_ids: torch.Tensor, local_num_experts: int): + reorder_topk_ids, reorder_ids = torch.sort(local_topk_ids.view(-1), stable=True) + + seg_indptr = torch.zeros( + local_num_experts + 1, device=local_topk_ids.device, dtype=torch.int64 + ) + src2dst = torch.empty( + local_topk_ids.numel(), device=local_topk_ids.device, dtype=torch.int32 + ) + + BLOCK_SIZE = 512 + grid = (triton.cdiv(local_topk_ids.numel(), BLOCK_SIZE),) + compute_src2dst_triton_kernel[grid]( + reorder_ids, src2dst, local_topk_ids.numel(), BLOCK_SIZE + ) + + return reorder_topk_ids, src2dst, seg_indptr + + +@triton.jit +def pre_reorder_triton_kernel_for_cutlass_moe( + input_ptr, + gateup_input_ptr, + src2dst_ptr, + topk_ids_ptr, + a1_scales_ptr, + num_experts, + topk, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + OutDtype = gateup_input_ptr.dtype.element_ty + + src_idx = tl.program_id(0) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + + src_ptr = input_ptr + src_idx * hidden_size + for idx in range(topk): + expert_id = tl.load(topk_ids_ptr + idx) + if expert_id != num_experts: + if a1_scales_ptr is not None: + scale = 1.0 / tl.load(a1_scales_ptr) + else: + scale = 1.0 + + dst_idx = tl.load(src2dst_ptr + idx) + dst_ptr = gateup_input_ptr + dst_idx * hidden_size + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + in_data = tl.load(src_ptr + offset, mask=mask).to(tl.float32) + out_data = (in_data * scale).to(OutDtype) + tl.store(dst_ptr + offset, out_data, mask=mask) + + +@triton.jit +def pre_reorder_triton_kernel( + input_ptr, + gateup_input_ptr, + src2dst_ptr, + topk_ids_ptr, + a1_scales_ptr, + start_expert_id, + end_expert_id, + topk, + hidden_size, + BLOCK_SIZE: tl.constexpr, + use_per_token_if_dynamic: tl.constexpr, +): + OutDtype = gateup_input_ptr.dtype.element_ty + + src_idx_int32 = tl.program_id(0) + src_idx = src_idx_int32.to(tl.int64) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + src_ptr = input_ptr + src_idx * hidden_size + + vec = tl.arange(0, BLOCK_SIZE) + + if a1_scales_ptr is not None and use_per_token_if_dynamic: + scale = 1.0 / tl.load(a1_scales_ptr + src_idx) + + for idx in range(topk): + expert_id = tl.load(topk_ids_ptr + idx) + if expert_id >= start_expert_id and expert_id <= end_expert_id: + if a1_scales_ptr is not None: + if not use_per_token_if_dynamic: + scale = 1.0 / tl.load(a1_scales_ptr + expert_id - start_expert_id) + else: + scale = 1.0 + + dst_idx_int32 = tl.load(src2dst_ptr + idx) + dst_idx = dst_idx_int32.to(tl.int64) + dst_ptr = gateup_input_ptr + dst_idx * hidden_size + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < hidden_size + in_data = tl.load(src_ptr + offset, mask=mask).to(tl.float32) + out_data = (in_data * scale).to(OutDtype) + tl.store(dst_ptr + offset, out_data, mask=mask) + + +@triton.jit +def silu_and_mul_triton_kernel( + gateup_output, + down_input, + hidden_size, + reorder_topk_ids, + scales, + start_expert_id, + end_expert_id, + BLOCK_SIZE: tl.constexpr, +): + InDtype = gateup_output.dtype.element_ty + OutDtype = down_input.dtype.element_ty + + half_hidden_size = hidden_size // 2 + + pid = tl.program_id(0) + expert_id = tl.load(reorder_topk_ids + pid) + if expert_id >= start_expert_id and expert_id <= end_expert_id: + gateup_output_ptr = gateup_output + pid * hidden_size + gate_output_ptr = gateup_output_ptr + up_output_ptr = gateup_output_ptr + half_hidden_size + down_input_ptr = down_input + pid * half_hidden_size + + if scales is not None: + scale = tl.load(scales + expert_id - start_expert_id) + scale = (1 / scale).to(InDtype) + else: + scale = 1 + + for start_offset in tl.range(0, half_hidden_size, BLOCK_SIZE): + offset = start_offset + tl.arange(0, BLOCK_SIZE) + mask = offset < half_hidden_size + + gate_output = tl.load(gate_output_ptr + offset, mask=mask).to(tl.float32) + up_output = tl.load(up_output_ptr + offset, mask=mask) + + # silu & mul & quantize + gate_output = gate_output * tl.sigmoid(gate_output) + gate_output = gate_output.to(InDtype) + + silu_mul_output = gate_output * up_output * scale + silu_mul_output = silu_mul_output.to(OutDtype) + tl.store(down_input_ptr + offset, silu_mul_output, mask=mask) + + +# copy from https://github.com/ModelTC/lightllm/blob/a000ab69098654df4731f5b12587dd4e7f0a4f41/lightllm/common/fused_moe/moe_silu_and_mul_mix_quant_ep.py +@triton.jit +def _silu_and_mul_post_quant_kernel( + input_ptr, + stride_input_0, + stride_input_1, + stride_input_2, + output_ptr, + stride_output_0, + stride_output_1, + stride_output_2, + output_scale_ptr, + stride_output_scale_0, + stride_output_scale_1, + stride_output_scale_2, + masked_m_ptr, + size_n, + fp8_max, + fp8_min, + BLOCK_N: tl.constexpr, + NUM_STAGE: tl.constexpr, + SCALE_UE8M0: tl.constexpr, +): + expert_id = tl.program_id(2) + token_id = tl.program_id(1) + hidden_dim_block_index = tl.program_id(0) + + block_num_per_expert = tl.num_programs(1) + + token_num_cur_expert = tl.load(masked_m_ptr + expert_id) + + stride_input_0 = tl.cast(stride_input_0, dtype=tl.int64) + stride_output_0 = tl.cast(stride_output_0, dtype=tl.int64) + stride_input_1 = tl.cast(stride_input_1, dtype=tl.int64) + stride_output_1 = tl.cast(stride_output_1, dtype=tl.int64) + + offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N) + input_ptr_offs = input_ptr + expert_id * stride_input_0 + offs_in_d + output_ptr_offs = output_ptr + expert_id * stride_output_0 + offs_in_d + output_scale_offs = ( + output_scale_ptr + + expert_id * stride_output_scale_0 + + hidden_dim_block_index * stride_output_scale_2 + ) + + for token_index in tl.range( + token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE + ): + gate = tl.load( + input_ptr_offs + token_index * stride_input_1, + mask=offs_in_d < size_n, + other=0.0, + ).to(tl.float32) + up = tl.load( + input_ptr_offs + token_index * stride_input_1 + size_n, + mask=offs_in_d < size_n, + other=0.0, + ) + gate = gate / (1 + tl.exp(-gate)) + gate = gate.to(input_ptr.dtype.element_ty) + gate_up = up * gate + _absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10) + output_s = _absmax / fp8_max + if SCALE_UE8M0: + output_s = tl.exp2(tl.ceil(tl.log2(tl.abs(output_s)))) + output_q = tl.clamp(gate_up / output_s, fp8_min, fp8_max).to( + output_ptr.dtype.element_ty + ) + tl.store( + output_ptr_offs + token_index * stride_output_1, + output_q, + mask=offs_in_d < size_n, + ) + tl.store( + output_scale_offs + token_index * stride_output_scale_1, + output_s, + ) + + +def silu_and_mul_masked_post_quant_fwd( + input: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor, + quant_group_size: int, + masked_m: torch.Tensor, + scale_ue8m0: bool = False, +): + """ + input shape [expert_num, token_num_padded, hidden_dim] + output shape [expert_num, token_num_padded, hidden_dim // 2], dtype fp8 + output_scale [expert_num token_num_paddded, hidden_dim // 2 // 128] dtype float32 + quant_group_size int, + masked_m shape [expert_num], + """ + + assert input.is_contiguous() + assert output.dtype == torch.float8_e4m3fn + assert output.is_contiguous() + assert len(input.shape) == 3 + assert input.shape[0] == masked_m.shape[0] + assert input.shape[-1] % 2 == 0 + + size_n = input.shape[-1] // 2 + assert size_n % quant_group_size == 0 + + expert_num = len(masked_m) + + if expert_num < 4: + BLOCK_NUM_PER_EXPERT = 64 + else: + BLOCK_NUM_PER_EXPERT = 32 + + BLOCK_N = quant_group_size + num_warps = 1 + NUM_STAGES = 6 + hidden_dim_split_block_num = triton.cdiv(size_n, BLOCK_N) + assert BLOCK_N % quant_group_size == 0 + + grid = ( + hidden_dim_split_block_num, + BLOCK_NUM_PER_EXPERT, + expert_num, + ) + + finfo = torch.finfo(torch.float8_e4m3fn) + fp8_max = finfo.max + fp8_min = -fp8_max + + _silu_and_mul_post_quant_kernel[grid]( + input, + *input.stride(), + output, + *output.stride(), + output_scale, + *output_scale.stride(), + masked_m, + size_n, + fp8_max, + fp8_min, + BLOCK_N=BLOCK_N, + NUM_STAGE=NUM_STAGES, + num_warps=num_warps, + SCALE_UE8M0=scale_ue8m0, + ) + return + + +@triton.jit +def tanh(x): + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def gelu_and_mul_triton_kernel( + gateup_output, + down_input, + hidden_size, + reorder_topk_ids, + scales, + start_expert_id, + end_expert_id, + BLOCK_SIZE: tl.constexpr, +): + InDtype = gateup_output.dtype.element_ty + OutDtype = down_input.dtype.element_ty + + half_hidden_size = hidden_size // 2 + + pid = tl.program_id(0) + expert_id = tl.load(reorder_topk_ids + pid) + if expert_id >= start_expert_id and expert_id <= end_expert_id: + gateup_output_ptr = gateup_output + pid * hidden_size + gate_output_ptr = gateup_output_ptr + up_output_ptr = gateup_output_ptr + half_hidden_size + down_input_ptr = down_input + pid * half_hidden_size + + if scales is not None: + scale = tl.load(scales + expert_id - start_expert_id) + scale = (1 / scale).to(InDtype) + else: + scale = 1 + + for start_offset in tl.range(0, half_hidden_size, BLOCK_SIZE): + offset = start_offset + tl.arange(0, BLOCK_SIZE) + mask = offset < half_hidden_size + + gate_output = tl.load(gate_output_ptr + offset, mask=mask).to(tl.float32) + up_output = tl.load(up_output_ptr + offset, mask=mask) + + # gelu & mul & quantize + # https://pytorch.org/docs/stable/generated/torch.nn.GELU.html + # sqrt(2/pi) + kAlpha = 0.7978845608028654 + gate_output = ( + 0.5 + * gate_output + * ( + 1 + + tanh( + kAlpha + * ( + gate_output + + 0.044715 * gate_output * gate_output * gate_output + ) + ) + ) + ) + gate_output = gate_output.to(InDtype) + + gelu_mul_output = gate_output * up_output * scale + gelu_mul_output = gelu_mul_output.to(OutDtype) + tl.store(down_input_ptr + offset, gelu_mul_output, mask=mask) + + +@triton.jit +def post_reorder_triton_kernel( + down_output_ptr, + output_ptr, + src2dst_ptr, + topk_ids_ptr, + topk_weights_ptr, + start_expert_id, + end_expert_id, + topk, + hidden_size, + dst_start, + BLOCK_SIZE: tl.constexpr, +): + InDtype = down_output_ptr.dtype.element_ty + + src_idx_int32 = tl.program_id(0) + src_idx = src_idx_int32.to(tl.int64) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + topk_weights_ptr = topk_weights_ptr + src_idx * topk + + computed = False + store_ptr = output_ptr + src_idx * hidden_size + + vec = tl.arange(0, BLOCK_SIZE) + + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < hidden_size + + sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype) + for idx in range(topk): + expert_id = tl.load(topk_ids_ptr + idx) + if expert_id >= start_expert_id and expert_id <= end_expert_id: + computed = True + dst_idx_int32 = tl.load(src2dst_ptr + idx) + dst_idx = dst_idx_int32.to(tl.int64) + dst_idx = dst_idx - dst_start + weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype) + load_ptr = down_output_ptr + dst_idx * hidden_size + in_data = tl.load(load_ptr + offset, mask=mask) + sum_vec += in_data * weigh_scale + tl.store(store_ptr + offset, sum_vec, mask=mask) + + if computed == False: + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < hidden_size + tl.store( + store_ptr + offset, tl.zeros([BLOCK_SIZE], dtype=InDtype), mask=mask + ) + + +@triton.jit +def post_reorder_triton_kernel_for_cutlass_moe( + down_output_ptr, + output_ptr, + src2dst_ptr, + topk_ids_ptr, + topk_weights_ptr, + num_experts, + topk, + hidden_size, + dst_start, + BLOCK_SIZE: tl.constexpr, +): + InDtype = down_output_ptr.dtype.element_ty + + src_idx_int32 = tl.program_id(0) + src_idx = src_idx_int32.to(tl.int64) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + topk_weights_ptr = topk_weights_ptr + src_idx * topk + + store_ptr = output_ptr + src_idx * hidden_size + + vec = tl.arange(0, BLOCK_SIZE) + + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < hidden_size + + sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype) + for idx in range(topk): + expert_id = tl.load(topk_ids_ptr + idx) + if expert_id != num_experts: + dst_idx_int32 = tl.load(src2dst_ptr + idx) + dst_idx = dst_idx_int32.to(tl.int64) + dst_idx = dst_idx - dst_start + weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype) + load_ptr = down_output_ptr + dst_idx * hidden_size + in_data = tl.load(load_ptr + offset, mask=mask) + sum_vec += in_data * weigh_scale + tl.store(store_ptr + offset, sum_vec, mask=mask) + + +@triton.jit +def compute_m_range( + pid, + batch_size, + seg_indptr, + weight_indices, + m_num_tiles_indptr, + BLOCK_SIZE_M: tl.constexpr, +): + idx = 0 + for bs in range(batch_size): + tiles = tl.load(m_num_tiles_indptr + bs) + if pid >= tiles: + idx = bs + + idx_start = tl.load(m_num_tiles_indptr + idx) + + m_range_start = tl.load(seg_indptr + idx) + (pid - idx_start) * BLOCK_SIZE_M + m_range_end = min(tl.load(seg_indptr + idx + 1), m_range_start + BLOCK_SIZE_M) + expert_id = tl.load(weight_indices + idx) + return m_range_start, m_range_end, expert_id + + +@triton.jit +def grouped_gemm_triton_kernel( + a, + b, + c, + batch_size, + N, + K, + seg_indptr, + weight_indices, + m_num_tiles_indptr, + scale_a, + scale_b, + use_fp8_w8a8: tl.constexpr, + group_n: tl.constexpr, + group_k: tl.constexpr, + a_stride_0: tl.constexpr, + b_stride_0: tl.constexpr, + b_stride_1: tl.constexpr, + as_stride_0: tl.constexpr, + as_stride_1: tl.constexpr, + bs_stride_0: tl.constexpr, + bs_stride_2: tl.constexpr, + bs_stride_1: tl.constexpr, + use_per_token_if_dynamic: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + c_dtype = c.dtype.element_ty + + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + total_m_block = tl.load(m_num_tiles_indptr + batch_size) + if pid_m >= total_m_block: + return + + m_range_start, m_range_end, expert_id = compute_m_range( + pid_m, batch_size, seg_indptr, weight_indices, m_num_tiles_indptr, BLOCK_SIZE_M + ) + if m_range_end - m_range_start == 0: + return + + n_range_start = pid_n * BLOCK_SIZE_N + n_range_end = min(n_range_start + BLOCK_SIZE_N, N) + + offs_am = tl.arange(0, BLOCK_SIZE_M) + offs_bn = tl.arange(0, BLOCK_SIZE_N) + + offs_am = tl.where(offs_am < m_range_end - m_range_start, offs_am, 0) + offs_bn = tl.where(offs_bn < n_range_end - n_range_start, offs_bn, 0) + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + + a_ptr = a + (m_range_start + offs_am[:, None]) * a_stride_0 + offs_k[None, :] + b_ptr = b + ( + (expert_id * b_stride_0) + + (n_range_start + offs_bn[:, None]) * b_stride_1 + + offs_k[None, :] + ) + + if group_k > 0 and group_n > 0: + a_scale_ptrs = scale_a + (m_range_start + offs_am[:, None]) * as_stride_0 + offs_bsn = (n_range_start + offs_bn) // group_n + b_scale_ptrs = scale_b + (expert_id * bs_stride_0) + offs_bsn * bs_stride_1 + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a_tile = tl.load( + a_ptr, mask=offs_k[None, :] < (K - k * BLOCK_SIZE_K), other=0.0 + ) + b_tile = tl.load( + b_ptr, mask=offs_k[None, :] < (K - k * BLOCK_SIZE_K), other=0.0 + ) + + if group_k > 0 and group_n > 0: + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_scale = tl.load(a_scale_ptrs + offs_ks * as_stride_1) + b_scale = tl.load(b_scale_ptrs + offs_ks * bs_stride_2) + accumulator += tl.dot(a_tile, b_tile.T) * a_scale * b_scale[None, :] + else: + accumulator = tl.dot(a_tile, b_tile.T, accumulator) + a_ptr += BLOCK_SIZE_K + b_ptr += BLOCK_SIZE_K + + if use_fp8_w8a8 and not (group_k > 0 and group_n > 0): + if use_per_token_if_dynamic: + scale_a_value = tl.load(scale_a + (m_range_start + offs_am[:, None])) + else: + scale_a_value = tl.load(scale_a + expert_id) + scale_b_value = tl.load(scale_b + expert_id) + accumulator *= scale_a_value * scale_b_value + + c_tile = accumulator.to(c_dtype) + + offs_cm = m_range_start + tl.arange(0, BLOCK_SIZE_M) + offs_cn = n_range_start + tl.arange(0, BLOCK_SIZE_N) + c_ptr = c + offs_cm[:, None] * N + offs_cn[None, :] + c_mask = (offs_cm[:, None] < m_range_end) & (offs_cn[None, :] < n_range_end) + tl.store(c_ptr, c_tile, mask=c_mask) + + +@triton.jit +def compute_m_num_tiles_indptr( + m_num_tiles_indptr, seg_indptr, batch_size: tl.constexpr, BLOCK_SIZE_M: tl.constexpr +): + for bs in range(batch_size): + m = tl.load(seg_indptr + bs + 1) - tl.load(seg_indptr + bs) + cur_num_tiles = tl.cdiv(m, BLOCK_SIZE_M) + pre_num_tiles = tl.load(m_num_tiles_indptr + bs) + tl.store(m_num_tiles_indptr + bs + 1, pre_num_tiles + cur_num_tiles) + + +def grouped_gemm_triton( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + batch_size: int, + weight_column_major: bool, + seg_indptr: Optional[torch.Tensor] = None, + weight_indices: Optional[torch.Tensor] = None, + use_fp8_w8a8: bool = False, + scale_a: torch.Tensor = None, + scale_b: torch.Tensor = None, + block_shape: Optional[List[int]] = None, + c_dtype=None, + use_per_token_if_dynamic: bool = True, +): + assert weight_column_major == True # TODO: more + if use_fp8_w8a8 and block_shape is None: + assert scale_a is not None and scale_b is not None + + if block_shape is not None: + a_original = a + + assert len(block_shape) == 2 + block_n, block_k = block_shape[0], block_shape[1] + a, scale_a = per_token_group_quant_fp8(a, block_k) + + assert triton.cdiv(a.shape[-1], block_k) == scale_a.shape[-1] + assert triton.cdiv(b.shape[-2], block_n) == scale_b.shape[-2] + assert triton.cdiv(b.shape[-1], block_k) == scale_b.shape[-1] + + dispose_tensor(a_original) + + # TODO: adjust config or tune kernel + # Reduce block size to prevent L40 shared memory overflow. + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + } + + m_num_tiles_indptr = torch.zeros(batch_size + 1, device=a.device, dtype=torch.int64) + compute_m_num_tiles_indptr[(1,)]( + m_num_tiles_indptr, seg_indptr, batch_size, config["BLOCK_SIZE_M"] + ) + + if c is None: + assert c_dtype is not None + c = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=c_dtype) + + grid = lambda META: ( + triton.cdiv(a.size(0), META["BLOCK_SIZE_M"]) + batch_size, + triton.cdiv(b.size(1), META["BLOCK_SIZE_N"]), + ) + + if use_fp8_w8a8 and block_shape is None and use_per_token_if_dynamic: + assert ( + scale_a.shape[0] == a.shape[0] + ), f"scale_a.shape: {scale_a.shape}, a.shape: {a.shape}" + + grouped_gemm_triton_kernel[grid]( + a, + b, + c, + batch_size, + b.size(1), + b.size(2), + seg_indptr, + weight_indices, + m_num_tiles_indptr, + scale_a, + scale_b, + use_fp8_w8a8, + 0 if block_shape is None else block_shape[0], + 0 if block_shape is None else block_shape[1], + a.stride(0), + b.stride(0), + b.stride(1), + scale_a.stride(0) if scale_a is not None and scale_a.ndim == 2 else 0, + scale_a.stride(1) if scale_a is not None and scale_a.ndim == 2 else 0, + scale_b.stride(0) if scale_b is not None and scale_b.ndim >= 2 else 0, + scale_b.stride(2) if scale_b is not None and scale_b.ndim == 3 else 0, + scale_b.stride(1) if scale_b is not None and scale_b.ndim >= 2 else 0, + use_per_token_if_dynamic, + **config, + ) + return c + + +@triton.jit +def _fwd_kernel_ep_scatter_1( + num_recv_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts: tl.constexpr, + BLOCK_E: tl.constexpr, + BLOCK_EXPERT_NUM: tl.constexpr, +): + cur_expert = tl.program_id(0) + + offset_cumsum = tl.arange(0, BLOCK_EXPERT_NUM) + tokens_per_expert = tl.load( + num_recv_tokens_per_expert + offset_cumsum, + mask=offset_cumsum < num_experts, + other=0, + ) + cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert + tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts) + + cur_expert_start = tl.load(expert_start_loc + cur_expert) + cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert) + + m_indices_start_ptr = m_indices + cur_expert_start + off_expert = tl.arange(0, BLOCK_E) + + for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4): + tl.store( + m_indices_start_ptr + start_m + off_expert, + cur_expert, + ) + + +@triton.jit +def _fwd_kernel_ep_scatter_2( + total_token_num, + expert_start_loc, + recv_x, + recv_x_stride0, + recv_x_stride1, + recv_x_scale, + recv_x_scale_stride0, + recv_x_scale_stride1, + recv_topk, + recv_topk_stride0, + recv_topk_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + output_tensor_scale, + output_tensor_scale_stride0, + output_tensor_scale_stride1, + output_index, + output_index_stride0, + output_index_stride1, + topk_num: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + HIDDEN_SIZE_PAD: tl.constexpr, + SCALE_HIDDEN_SIZE: tl.constexpr, + SCALE_HIDDEN_SIZE_PAD: tl.constexpr, +): + start_token_id = tl.program_id(0) + grid_num = tl.num_programs(0) + + offset_in = tl.arange(0, HIDDEN_SIZE_PAD) + mask = offset_in < HIDDEN_SIZE + + index_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD) + mask_s = index_in_s < SCALE_HIDDEN_SIZE + + for token_id_int32 in range(start_token_id, total_token_num, grid_num): + token_id = token_id_int32.to(tl.int64) + to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask) + to_copy_s = tl.load( + recv_x_scale + + token_id * recv_x_scale_stride0 + + index_in_s * recv_x_scale_stride1, + mask=mask_s, + ) + + for topk_idx_int32 in tl.range(0, topk_num, 1, num_stages=4): + topk_index = topk_idx_int32.to(tl.int64) + expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index) + if expert_id >= 0: + dest_token_index_int32 = tl.atomic_add(expert_start_loc + expert_id, 1) + dest_token_index = dest_token_index_int32.to(tl.int64) + + tl.store( + output_index + token_id * output_index_stride0 + topk_index, + dest_token_index_int32, + ) + output_tensor_ptr = ( + output_tensor + dest_token_index * output_tensor_stride0 + ) + output_tensor_scale_ptr = ( + output_tensor_scale + dest_token_index * output_tensor_scale_stride0 + ) + tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask) + tl.store( + output_tensor_scale_ptr + index_in_s * output_tensor_scale_stride1, + to_copy_s, + mask=mask_s, + ) + + +# copy from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/deepep_scatter_gather.py +@torch.no_grad() +def ep_scatter( + recv_x: torch.Tensor, + recv_x_scale: torch.Tensor, + recv_topk: torch.Tensor, + num_recv_tokens_per_expert: torch.Tensor, + expert_start_loc: torch.Tensor, + output_tensor: torch.Tensor, + output_tensor_scale: torch.Tensor, + m_indices: torch.Tensor, + output_index: torch.Tensor, + scale_ue8m0: bool = False, +): + BLOCK_E = 128 # token num of per expert is aligned to 128 + BLOCK_D = 128 # block size of quantization + num_warps = 8 + num_experts = num_recv_tokens_per_expert.shape[0] + hidden_size = recv_x.shape[1] + # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts) + grid = num_experts + + scale_hidden_size = hidden_size // BLOCK_D + if scale_ue8m0: + # ue8m0 scales are packed here (4 scales per int32), + # hence the effective size of this dimension is divided by 4. + scale_hidden_size = ceil_div(scale_hidden_size, 4) + + assert m_indices.shape[0] % BLOCK_E == 0 + assert recv_x_scale.dtype == output_tensor_scale.dtype + assert recv_x_scale.shape[1] == output_tensor_scale.shape[1] == scale_hidden_size + + _fwd_kernel_ep_scatter_1[(grid,)]( + num_recv_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts=num_experts, + num_warps=num_warps, + BLOCK_E=BLOCK_E, + BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts), + ) + + grid = min(recv_topk.shape[0], 1024 * 8) + + _fwd_kernel_ep_scatter_2[(grid,)]( + recv_topk.shape[0], + expert_start_loc, + recv_x, + recv_x.stride(0), + recv_x.stride(1), + recv_x_scale, + recv_x_scale.stride(0), + recv_x_scale.stride(1), + recv_topk, + recv_topk.stride(0), + recv_topk.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + output_tensor_scale, + output_tensor_scale.stride(0), + output_tensor_scale.stride(1), + output_index, + output_index.stride(0), + output_index.stride(1), + topk_num=recv_topk.shape[1], + num_warps=num_warps, + HIDDEN_SIZE=hidden_size, + HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size), + SCALE_HIDDEN_SIZE=scale_hidden_size, + SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size), + ) + return + + +@triton.jit +def _fwd_kernel_ep_gather( + total_token_num, + input_tensor, + input_tensor_stride0, + input_tensor_stride1, + recv_topk_ids, + recv_topk_ids_stride0, + recv_topk_ids_stride1, + recv_topk_weight, + recv_topk_weight_stride0, + recv_topk_weight_stride1, + input_index, + input_index_stride0, + input_index_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + topk_num: tl.constexpr, + BLOCK_D: tl.constexpr, +): + cur_block_int32 = tl.program_id(0) + cur_block = cur_block_int32.to(tl.int64) + + start_cur_token_int32 = tl.program_id(1) + + grid_num = tl.num_programs(1) + + for cur_token_int32 in range(start_cur_token_int32, total_token_num, grid_num): + cur_token = cur_token_int32.to(tl.int64) + + off_d = tl.arange(0, BLOCK_D) + accumulator = tl.zeros([BLOCK_D], dtype=tl.float32) + + for topk_index_int32 in range(0, topk_num): + topk_index = topk_index_int32.to(tl.int64) + + expert_id = tl.load( + recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index + ) + if expert_id >= 0: + source_token_index_int32 = tl.load( + input_index + cur_token * input_index_stride0 + topk_index + ) + source_token_index = source_token_index_int32.to(tl.int64) + + acc_weight = tl.load( + recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index + ) + tmp = tl.load( + input_tensor + + source_token_index * input_tensor_stride0 + + cur_block * BLOCK_D + + off_d + ) + accumulator += tmp.to(tl.float32) * acc_weight + + tl.store( + output_tensor + + cur_token * output_tensor_stride0 + + cur_block * BLOCK_D + + off_d, + accumulator.to(output_tensor.dtype.element_ty), + ) + + +@torch.no_grad() +def ep_gather( + input_tensor: torch.Tensor, + recv_topk_ids: torch.Tensor, + recv_topk_weight: torch.Tensor, + input_index: torch.Tensor, + output_tensor: torch.Tensor, +): + BLOCK_D = 1024 if not is_in_ci() else 128 # block size of quantization + num_warps = 2 + num_tokens = output_tensor.shape[0] + hidden_size = input_tensor.shape[1] + assert hidden_size % BLOCK_D == 0 + grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024)) + _fwd_kernel_ep_gather[grid]( + num_tokens, + input_tensor, + input_tensor.stride(0), + input_tensor.stride(1), + recv_topk_ids, + recv_topk_ids.stride(0), + recv_topk_ids.stride(1), + recv_topk_weight, + recv_topk_weight.stride(0), + recv_topk_weight.stride(1), + input_index, + input_index.stride(0), + input_index.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + topk_num=recv_topk_ids.shape[1], + num_warps=num_warps, + BLOCK_D=BLOCK_D, + ) + return + + +# copy from +# https://github.com/deepseek-ai/DeepGEMM/blob/bd2a77552886b98c205af12f8d7d2d61247c4b27/deep_gemm/jit_kernels/utils.py#L58 +def get_tma_aligned_size(x: int, element_size: int) -> int: + """ + Global memory address of TMA must be 16-byte aligned. + Since we use column-major layout for the LHS scaling tensor, + the M-axis of the LHS scaling tensor needs to be padded to a multiple of 16 bytes. + + Arguments: + x: original M-axis shape of the LHS scaling tensor. + element_size: element size of the LHS scaling tensor. + + Returns: + M-axis shape of the LHS scaling tensor after padding. + """ + tma_alignment_bytes = 16 + assert tma_alignment_bytes % element_size == 0 + alignment = tma_alignment_bytes // element_size + return ceil_div(x, alignment) * alignment + + +@triton.jit +def _tma_align_input_scale_kernel( + input_scale_ptr, + output_ptr, + m, + k_div_block_size, + input_scale_stride_m, + input_scale_stride_k, + output_stride_m, + output_stride_k, + BLOCK_SIZE_K: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + grid_m = tl.num_programs(0) + k_offsets = tl.arange(0, BLOCK_SIZE_K) + + for m_base in range(pid_m, m, grid_m): + input_offset = ( + input_scale_ptr + + m_base * input_scale_stride_m + + k_offsets * input_scale_stride_k + ) + input_data = tl.load(input_offset, mask=k_offsets < k_div_block_size) + + output_offset = ( + output_ptr + k_offsets * output_stride_k + m_base * output_stride_m + ) + tl.store(output_offset, input_data, mask=k_offsets < k_div_block_size) + + +# copy from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/quantization/triton_quant/fp8/fp8act_quant_kernel.py +def tma_align_input_scale(input_scale: torch.Tensor): + assert input_scale.dim() == 2 + m, k_div_block_size = input_scale.shape + padd_m = get_tma_aligned_size(m, input_scale.element_size()) + output = torch.empty( + (k_div_block_size, padd_m), dtype=input_scale.dtype, device=input_scale.device + ) + + grid_m = min(m, 8192) + BLOCK_SIZE_K = triton.next_power_of_2(k_div_block_size) + + _tma_align_input_scale_kernel[(grid_m,)]( + input_scale_ptr=input_scale, + output_ptr=output, + m=m, + k_div_block_size=k_div_block_size, + input_scale_stride_m=input_scale.stride(0), + input_scale_stride_k=input_scale.stride(1), + output_stride_m=output.stride(1), # Note: these are swapped + output_stride_k=output.stride(0), # for column-major + BLOCK_SIZE_K=BLOCK_SIZE_K, + ) + return output.t()[:m] + + +@triton.jit +def compute_masked_m_triton_kernel(seg_indptr, masked_m): + expert_id = tl.program_id(0) + start = tl.load(seg_indptr + expert_id) + end = tl.load(seg_indptr + expert_id + 1) + tl.store(masked_m + expert_id, (end - start)) + + +@triton.jit +def deepgemm_compute_src2dst_triton_kernel( + topk_ids, + reorder_ids, + seg_indptr, + src2dst, + m_max, + num_toks, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = dst_id < num_toks + src_id = tl.load(reorder_ids + dst_id, mask=mask) + expert_id = tl.load(topk_ids + src_id, mask=(src_id < num_toks)) + expert_dst_start = tl.load(seg_indptr + expert_id) + expert_dst_offset = dst_id - expert_dst_start + dst_id = expert_id * m_max + expert_dst_offset + tl.store(src2dst + src_id, dst_id, mask=mask) + + +@triton.jit +def fill_gateup_input_triton_kernel( + input_ptr, + scale_ptr, + gateup_input_ptr, + gateup_input_scale_ptr, + src2dst_ptr, + topk_ids_ptr, + start_expert_id, + end_expert_id, + topk, + m_max, + hidden_size, + scale_size, + BLOCK_SIZE: tl.constexpr, +): + + src_idx_int32 = tl.program_id(0) + src_idx = src_idx_int32.to(tl.int64) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + src_ptr = input_ptr + src_idx * hidden_size + scale_src_ptr = scale_ptr + src_idx * scale_size + + vec = tl.arange(0, BLOCK_SIZE) + for idx in range(topk): + expert_id = tl.load(topk_ids_ptr + idx) + if expert_id >= start_expert_id and expert_id <= end_expert_id: + dst_idx_int32 = tl.load(src2dst_ptr + idx) + dst_idx = dst_idx_int32.to(tl.int64) + dst_idx = dst_idx - start_expert_id * m_max + dst_ptr = gateup_input_ptr + dst_idx * hidden_size + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < hidden_size + in_data = tl.load(src_ptr + offset, mask=mask) + tl.store(dst_ptr + offset, in_data, mask=mask) + scale_dst_ptr = gateup_input_scale_ptr + dst_idx * scale_size + for start_offset in tl.range(0, scale_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < scale_size + in_scale = tl.load(scale_src_ptr + offset, mask=mask) + tl.store(scale_dst_ptr + offset, in_scale, mask=mask) + + +def moe_ep_deepgemm_preprocess( + topk_ids: torch.Tensor, + num_experts: int, + hidden_states: torch.Tensor, + top_k: int, + start_expert_id, + end_expert_id, + block_shape, + output_dtype: torch.dtype = torch.float8_e4m3fn, +): + reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True) + seg_indptr = torch.zeros(num_experts + 1, device=topk_ids.device, dtype=torch.int64) + src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int32) + masked_m = torch.zeros(num_experts, device=topk_ids.device, dtype=torch.int32) + + compute_seg_indptr_triton_kernel[(num_experts,)]( + reorder_topk_ids, seg_indptr, topk_ids.numel() + ) + + grid = lambda meta: (triton.cdiv(topk_ids.numel(), meta["BLOCK_SIZE"]),) + compute_masked_m_triton_kernel[(num_experts,)](seg_indptr, masked_m) + + # For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m}) https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/jit_kernels/m_grouped_gemm.py#L165 + m_max = (hidden_states.size(0) + 255) // 256 * 256 + expected_m = (topk_ids.numel() + num_experts - 1) // num_experts + gateup_input = torch.empty( + (int(end_expert_id - start_expert_id + 1), m_max, hidden_states.size(1)), + device=hidden_states.device, + dtype=output_dtype, + ) + + deepgemm_compute_src2dst_triton_kernel[grid]( + topk_ids, + reorder_ids, + seg_indptr, + src2dst, + m_max, + topk_ids.numel(), + BLOCK_SIZE=256, + ) + + if block_shape is None: + block_shape = [128, 128] + assert len(block_shape) == 2 + block_n, block_k = block_shape[0], block_shape[1] + hidden_states, scale = per_token_group_quant_fp8(hidden_states, block_k) + + gateup_input_scale = torch.empty( + (gateup_input.size(0), gateup_input.size(1), scale.size(1)), + device=hidden_states.device, + dtype=scale.dtype, + ) + + fill_gateup_input_triton_kernel[(hidden_states.shape[0],)]( + hidden_states, + scale, + gateup_input, + gateup_input_scale, + src2dst, + topk_ids, + start_expert_id, + end_expert_id, + top_k, + m_max, + hidden_states.size(1), + scale.size(1), + BLOCK_SIZE=1024, + ) + + return ( + m_max, + masked_m[start_expert_id : (end_expert_id + 1)], + expected_m, + src2dst, + gateup_input, + gateup_input_scale, + ) + + +@triton.jit +def compute_identity_kernel( + top_k, + hidden_states_ptr, + expert_scales_ptr, + num_tokens, + output_ptr, + hidden_dim, + scales_stride, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + + batch_id = pid // (hidden_dim // BLOCK_SIZE) + dim_offset = pid % (hidden_dim // BLOCK_SIZE) * BLOCK_SIZE + + if batch_id >= num_tokens or dim_offset >= hidden_dim: + return + + h = tl.load( + hidden_states_ptr + + batch_id * hidden_dim + + dim_offset + + tl.arange(0, BLOCK_SIZE), + mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, + ) + + result = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for i in range(top_k): + scale = tl.load(expert_scales_ptr + batch_id * scales_stride + i) + result += h * scale + + tl.store( + output_ptr + batch_id * hidden_dim + dim_offset + tl.arange(0, BLOCK_SIZE), + result, + mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, + ) + + +def zero_experts_compute_triton( + expert_indices, expert_scales, num_experts, zero_expert_type, hidden_states +): + N = expert_indices.numel() + top_k = expert_indices.size(-1) + grid = lambda meta: (triton.cdiv(N, meta["BLOCK_SIZE"]),) + + if zero_expert_type == "identity": + zero_expert_mask = expert_indices < num_experts + zero_expert_scales = expert_scales.clone() + zero_expert_scales[zero_expert_mask] = 0.0 + + normal_expert_mask = expert_indices >= num_experts + expert_indices[normal_expert_mask] = -1 + expert_scales[normal_expert_mask] = 0.0 + + output = torch.zeros_like(hidden_states).to(hidden_states.device) + hidden_dim = hidden_states.size(-1) + num_tokens = hidden_states.size(0) + + grid = lambda meta: (num_tokens * (hidden_dim // meta["BLOCK_SIZE"]),) + compute_identity_kernel[grid]( + top_k, + hidden_states, + zero_expert_scales, + num_tokens, + output, + hidden_dim, + zero_expert_scales.stride(0), + BLOCK_SIZE=256, + ) + + return output diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py new file mode 100644 index 000000000..ef33665c3 --- /dev/null +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -0,0 +1,869 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Optional, Union + +import torch + +from sglang.srt.distributed.parallel_state import get_moe_expert_parallel_world_size +from sglang.srt.layers.moe import ( + get_deepep_mode, + get_moe_a2a_backend, + get_moe_runner_backend, + should_use_flashinfer_trtllm_moe, +) +from sglang.srt.layers.moe.ep_moe.kernels import ( + ep_gather, + ep_scatter, + moe_ep_deepgemm_preprocess, + post_reorder_triton_kernel, + silu_and_mul_masked_post_quant_fwd, + tma_align_input_scale, +) +from sglang.srt.layers.moe.fused_moe_triton.layer import FlashInferFusedMoE, FusedMoE +from sglang.srt.layers.moe.topk import TopKOutput +from sglang.srt.layers.quantization import deep_gemm_wrapper +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.fp8 import Fp8Config +from sglang.srt.layers.quantization.fp8_kernel import ( + is_fp8_fnuz, + sglang_per_token_group_quant_fp8, +) +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.utils import ceil_div, dispose_tensor, get_bool_env_var, is_hip, is_npu + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + DeepEPLLOutput, + DeepEPNormalOutput, + DispatchOutput, + ) + +_is_hip = is_hip() +_is_npu = is_npu() +_is_fp8_fnuz = is_fp8_fnuz() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if not (_is_npu or _is_hip): + from sgl_kernel import silu_and_mul + +if _use_aiter: + from aiter import ActivationType, QuantType + from aiter.fused_moe import fused_moe + +logger = logging.getLogger(__name__) + + +# TODO(kaixih@nvidia): ideally we should merge this logic into +# `fill_gateup_input_triton_kernel` to directly generate e8m0 scale. +@torch.compile +def _cast_to_e8m0_with_rounding_up(x: torch.Tensor) -> torch.Tensor: + temp = x.to(torch.float32).view(torch.int32) + exp = torch.bitwise_right_shift(temp, 23) + mant = torch.bitwise_and(temp, 0x7FFFFF) + is_ru = torch.logical_and( + torch.logical_and((mant > 0), (exp != 0xFE)), + ~torch.logical_and((exp == 0), (mant <= 0x400000)), + ) + exp = torch.where(is_ru, exp + 1, exp) + new_x = exp.to(torch.uint8).view(torch.int) + return new_x.transpose(1, 2).contiguous().transpose(1, 2) + + +class EPMoE(FusedMoE): + """ + MoE Expert Parallel Impl + + + """ + + def __init__( + self, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + layer_id: int, + num_fused_shared_experts: int = 0, + params_dtype: Optional[torch.dtype] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + activation: str = "silu", + routed_scaling_factor: Optional[float] = None, + gemm1_alpha: Optional[float] = None, + gemm1_clamp_limit: Optional[float] = None, + with_bias: bool = False, + ): + super().__init__( + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_fused_shared_experts=num_fused_shared_experts, + layer_id=layer_id, + top_k=top_k, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + activation=activation, + # apply_router_weight_on_input=apply_router_weight_on_input, + routed_scaling_factor=routed_scaling_factor, + gemm1_alpha=gemm1_alpha, + gemm1_clamp_limit=gemm1_clamp_limit, + with_bias=with_bias, + ) + + self.intermediate_size = intermediate_size + + if isinstance(quant_config, Fp8Config): + self.use_block_quant = getattr(self.quant_method, "block_quant", False) + self.block_shape = ( + self.quant_method.quant_config.weight_block_size + if self.use_block_quant + else None + ) + self.use_fp8_w8a8 = True + self.fp8_dtype = torch.float8_e4m3fn + self.activation_scheme = quant_config.activation_scheme + else: + self.use_fp8_w8a8 = False + self.use_block_quant = False + self.block_shape = None + self.activation_scheme = None + + def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput): + if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and self.use_fp8_w8a8: + return self.forward_deepgemm(hidden_states, topk_output) + else: + return super().forward(hidden_states, topk_output) + + def forward_deepgemm( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + ): + + self.w13_weight_fp8 = ( + self.w13_weight, + ( + self.w13_weight_scale_inv + if self.use_block_quant + else self.w13_weight_scale + ), + ) + self.w2_weight_fp8 = ( + self.w2_weight, + self.w2_weight_scale_inv if self.use_block_quant else self.w2_weight_scale, + ) + + assert self.quant_method is not None + assert self.moe_runner_config.activation == "silu" + + hidden_states_shape = hidden_states.shape + hidden_states_dtype = hidden_states.dtype + hidden_states_device = hidden_states.device + + topk_weights, topk_ids, _ = topk_output + + if not self.use_block_quant: + # Convert per-tensor quant to per-block quant by repeating scales for forward_deepgemm + scale_block_size = 128 + w13_weight_scale_n = 2 * ( + (self.intermediate_size + scale_block_size - 1) // scale_block_size + ) + w13_weight_scale_k = ( + hidden_states_shape[-1] + scale_block_size - 1 + ) // scale_block_size + w13_weight_scale = ( + self.w13_weight_scale.unsqueeze(1) + .repeat_interleave(w13_weight_scale_n, dim=1) + .unsqueeze(2) + .repeat_interleave(w13_weight_scale_k, dim=2) + ) + self.w13_weight_fp8 = ( + self.w13_weight, + w13_weight_scale, + ) + w2_weight_scale_n = ( + hidden_states_shape[-1] + scale_block_size - 1 + ) // scale_block_size + w2_weight_scale_k = ( + self.intermediate_size + scale_block_size - 1 + ) // scale_block_size + w2_weight_scale = ( + self.w2_weight_scale.unsqueeze(1) + .repeat_interleave(w2_weight_scale_n, dim=1) + .unsqueeze(2) + .repeat_interleave(w2_weight_scale_k, dim=2) + ) + self.w2_weight_fp8 = ( + self.w2_weight, + w2_weight_scale, + ) + + # PreReorder + m_max, masked_m, expected_m, src2dst, gateup_input, gateup_input_scale = ( + moe_ep_deepgemm_preprocess( + topk_ids, + self.num_experts, + hidden_states, + self.top_k, + self.start_expert_id, + self.end_expert_id, + self.block_shape, + ) + ) + + dispose_tensor(hidden_states) + + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + b, s_mn, s_k = gateup_input_scale.shape + assert ( + s_mn % 4 == 0 and s_k % 4 == 0 + ), f"scales must be aligned to 4, but got ({b}, {s_mn}, {s_k})" + + # GroupGemm-0 + gateup_input_fp8 = ( + gateup_input, + ( + _cast_to_e8m0_with_rounding_up(gateup_input_scale) + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + else deep_gemm_wrapper.get_mn_major_tma_aligned_tensor( + gateup_input_scale + ) + ), + ) + num_groups, m, k = gateup_input_fp8[0].size() + n = self.w13_weight.size(1) + gateup_output = torch.empty( + (num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16 + ) + deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked( + gateup_input_fp8, + self.w13_weight_fp8, + gateup_output, + masked_m, + expected_m, + ) + del gateup_input + del gateup_input_fp8 + + # Act + down_input = torch.empty( + ( + gateup_output.shape[0], + gateup_output.shape[1], + gateup_output.shape[2] // 2, + ), + device=hidden_states_device, + dtype=self.fp8_dtype, + ) + scale_block_size = 128 + down_input_scale = torch.empty( + ( + gateup_output.shape[0], + gateup_output.shape[1], + gateup_output.shape[2] // 2 // scale_block_size, + ), + device=hidden_states_device, + dtype=torch.float32, + ) + silu_and_mul_masked_post_quant_fwd( + gateup_output, + down_input, + down_input_scale, + scale_block_size, + masked_m, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + del gateup_output + + # GroupGemm-1 + n = self.w2_weight.size(1) + down_input_fp8 = ( + down_input, + ( + down_input_scale + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + else deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(down_input_scale) + ), + ) + down_output = torch.empty( + (num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16 + ) + deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked( + down_input_fp8, + self.w2_weight_fp8, + down_output, + masked_m, + expected_m, + ) + del down_input + del down_input_fp8 + + # PostReorder + output = torch.empty( + hidden_states_shape, dtype=hidden_states_dtype, device=hidden_states_device + ) + post_reorder_triton_kernel[(hidden_states_shape[0],)]( + down_output, + output, + src2dst, + topk_ids, + topk_weights, + self.start_expert_id, + self.end_expert_id, + self.top_k, + hidden_states_shape[1], + m_max * self.start_expert_id, + BLOCK_SIZE=512, + ) + if self.moe_runner_config.routed_scaling_factor is not None: + output *= self.moe_runner_config.routed_scaling_factor + return output + + +class DeepEPMoE(EPMoE): + """ + MoE Expert Parallel Impl based on DeepEP (https://github.com/deepseek-ai/DeepEP/tree/main) + """ + + _has_printed = False + + def __init__( + self, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + layer_id: int, + num_fused_shared_experts: int = 0, + params_dtype: Optional[torch.dtype] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + activation: str = "silu", + routed_scaling_factor: Optional[float] = None, + ): + super().__init__( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + layer_id=layer_id, + num_fused_shared_experts=num_fused_shared_experts, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + activation=activation, + routed_scaling_factor=routed_scaling_factor, + ) + self.deepep_mode = get_deepep_mode() + + # TODO: move to the beginning of the file + from sglang.srt.distributed.parallel_state import get_tp_group + from sglang.srt.two_batch_overlap import MaybeTboDeepEPDispatcher + + self.deepep_dispatcher = MaybeTboDeepEPDispatcher( + group=get_tp_group().device_group, + router_topk=self.top_k, + permute_fusion=True, + num_experts=self.num_experts, + num_local_experts=self.num_local_experts, + hidden_size=hidden_size, + params_dtype=params_dtype, + deepep_mode=self.deepep_mode, + async_finish=True, # TODO + return_recv_hook=True, + ) + + if self.deepep_mode.enable_low_latency() and not _is_npu: + # NPU supports low_latency deepep without deepgemm + assert ( + deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + ), f"DeepEP {self.deepep_mode} mode requires deep_gemm" + if _use_aiter: + # expert_mask is of size (self.num_local_experts + 1), + # the extra 1 is for invalid rank_id (in original deepep, the invalid rank_id is -1, but aiter does not allow -1, we use a mask to make those ids invalid) + # for instance, if we have 4 experts on this rank, we would have a expert_mask like: + # self.expert_mask = [1, 1, 1, 1, 0] + # idx from 0-3 is valid and will be processed, while idx == 4 will be masked out + self.expert_mask = torch.zeros( + (self.num_local_experts + 1), + device=torch.cuda.current_device(), + dtype=torch.int, + ) + # the last one is invalid rank_id + self.expert_mask[:-1] = 1 + elif not _is_npu: + self.w13_weight_fp8 = ( + self.w13_weight, + ( + self.w13_weight_scale_inv + if self.use_block_quant + else self.w13_weight_scale + ), + ) + self.w2_weight_fp8 = ( + self.w2_weight, + ( + self.w2_weight_scale_inv + if self.use_block_quant + else self.w2_weight_scale + ), + ) + + def forward( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + forward_batch: ForwardBatch, + ): + dispatch_output = self.dispatch( + hidden_states, topk_idx, topk_weights, forward_batch + ) + hidden_states = self.moe_impl(dispatch_output) + hidden_states = self.combine( + hidden_states, + dispatch_output.topk_idx, + dispatch_output.topk_weights, + forward_batch, + ) + return hidden_states + + def dispatch( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + forward_batch: ForwardBatch, + ): + return self.deepep_dispatcher.dispatch( + hidden_states=hidden_states, + topk_idx=topk_idx, + topk_weights=topk_weights, + forward_batch=forward_batch, + ) + + def moe_impl(self, dispatch_output: DispatchOutput): + from sglang.srt.layers.moe.token_dispatcher import DispatchOutputChecker + + if _use_aiter: + assert DispatchOutputChecker.format_is_deepep(dispatch_output) + # in forward_aiter, we skip token permutation and unpermutation, which have been fused inside aiter kernel + return self.forward_aiter(dispatch_output) + if _is_npu: + assert DispatchOutputChecker.format_is_deepep(dispatch_output) + return self.forward_npu(dispatch_output) + if DispatchOutputChecker.format_is_deepep_normal(dispatch_output): + assert deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and self.use_fp8_w8a8 + return self.forward_deepgemm_contiguous(dispatch_output) + elif DispatchOutputChecker.format_is_deepep_ll(dispatch_output): + assert deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and self.use_fp8_w8a8 + return self.forward_deepgemm_masked(dispatch_output) + else: + raise ValueError( + f"Dispatch output format {dispatch_output.format} is not supported" + ) + + def combine( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + forward_batch: ForwardBatch, + ): + return self.deepep_dispatcher.combine( + hidden_states=hidden_states, + topk_idx=topk_idx, + topk_weights=topk_weights, + forward_batch=forward_batch, + ) + + def forward_aiter( + self, + dispatch_output: Union[DeepEPNormalOutput, DeepEPLLOutput], + ): + hidden_states, topk_idx, topk_weights = ( + dispatch_output.hidden_states, + dispatch_output.topk_idx, + dispatch_output.topk_weights, + ) + if hidden_states.shape[0] == 0: + return hidden_states + # in original deepep, idx == -1 meaning invalid and will not be processed. + # aiter does not accept -1, we use a expert mask to make these idx invalid + # (idx == num_local_experts) meaning not used in aiter fused_moe + topk_idx_copy = topk_idx.to(torch.int32) + topk_idx_copy[topk_idx_copy == -1] = self.num_local_experts + + return fused_moe( + hidden_states, + self.w13_weight, + self.w2_weight, + topk_weights, + topk_idx_copy, + w1_scale=self.w13_weight_scale_inv, + w2_scale=self.w2_weight_scale_inv, + quant_type=QuantType.per_128x128, + activation=( + ActivationType.Silu + if self.moe_runner_config.activation == "silu" + else ActivationType.Gelu + ), + expert_mask=self.expert_mask, + ) + + def forward_deepgemm_contiguous( + self, + dispatch_output: DeepEPNormalOutput, + ): + hidden_states_fp8, topk_idx, topk_weights, num_recv_tokens_per_expert = ( + dispatch_output + ) + hidden_states_fp8, hidden_states_scale = hidden_states_fp8 + assert self.quant_method is not None + assert self.moe_runner_config.activation == "silu" + if num_recv_tokens_per_expert is None: + return hidden_states_fp8.bfloat16() + all_tokens = sum(num_recv_tokens_per_expert) + if all_tokens <= 0: + return hidden_states_fp8.bfloat16() + M, K = hidden_states_fp8.size() + N = self.w13_weight.size(1) + scale_block_size = 128 + + hidden_states_fp8_shape = hidden_states_fp8.shape + hidden_states_fp8_device = hidden_states_fp8.device + hidden_states_fp8_dtype = hidden_states_fp8.dtype + + input_tensor = [ + torch.empty( + (all_tokens, K), + device=hidden_states_fp8.device, + dtype=hidden_states_fp8.dtype, + ), + ( + # TODO check whether need `zeros` + torch.zeros( + (ceil_div(K // 128, 4), all_tokens), + device=hidden_states_fp8.device, + dtype=torch.int, + ).transpose(0, 1) + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + else torch.empty( + (all_tokens, K // 128), + device=hidden_states_fp8.device, + dtype=torch.float32, + ) + ), + ] + m_indices = torch.empty( + all_tokens, device=hidden_states_fp8.device, dtype=torch.int32 + ) + output_index = torch.empty_like(topk_idx) + + num_recv_tokens_per_expert_gpu = torch.tensor( + num_recv_tokens_per_expert, + dtype=torch.int32, + pin_memory=True, + device="cpu", + ).cuda(non_blocking=True) + expert_start_loc = torch.empty_like(num_recv_tokens_per_expert_gpu) + + ep_scatter( + hidden_states_fp8, + hidden_states_scale, + topk_idx, + num_recv_tokens_per_expert_gpu, + expert_start_loc, + input_tensor[0], + input_tensor[1], + m_indices, + output_index, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + dispose_tensor(hidden_states_fp8) + + gateup_output = torch.empty( + (all_tokens, N), + device=hidden_states_fp8_device, + dtype=torch.bfloat16, + ) + if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + input_tensor[1] = tma_align_input_scale(input_tensor[1]) + deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig( + input_tensor, self.w13_weight_fp8, gateup_output, m_indices + ) + del input_tensor + down_input = torch.empty( + ( + all_tokens, + N // 2, + ), + device=gateup_output.device, + dtype=torch.bfloat16, + ) + silu_and_mul(gateup_output.view(-1, N), down_input) + del gateup_output + down_output = torch.empty( + (all_tokens, K), + device=hidden_states_fp8_device, + dtype=torch.bfloat16, + ) + down_input_fp8, down_input_scale = sglang_per_token_group_quant_fp8( + down_input, + scale_block_size, + column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + del down_input + if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + down_input_scale = tma_align_input_scale(down_input_scale) + deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig( + (down_input_fp8, down_input_scale), + self.w2_weight_fp8, + down_output, + m_indices, + ) + del down_input_fp8, down_input_scale + + gather_out = torch.empty( + hidden_states_fp8_shape, + device=hidden_states_fp8_device, + dtype=torch.bfloat16, + ) + ep_gather(down_output, topk_idx, topk_weights, output_index, gather_out) + + return gather_out + + def forward_deepgemm_masked( + self, + dispatch_output: DeepEPLLOutput, + ): + hidden_states_fp8, _, _, masked_m, expected_m = dispatch_output + assert self.quant_method is not None + assert self.moe_runner_config.activation == "silu" + + # GroupGemm-0 + num_groups, m, k = hidden_states_fp8[0].size() + n = self.w13_weight.size(1) + expected_m = min(expected_m, m) + gateup_output = torch.empty( + (num_groups, m, n), device=hidden_states_fp8[0].device, dtype=torch.bfloat16 + ) + deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked( + hidden_states_fp8, + self.w13_weight_fp8, + gateup_output, + masked_m, + expected_m, + ) + dispose_tensor(hidden_states_fp8[0]) + + # Act + down_input = torch.empty( + ( + gateup_output.shape[0], + gateup_output.shape[1], + gateup_output.shape[2] // 2, + ), + device=gateup_output.device, + dtype=self.fp8_dtype, + ) + scale_block_size = 128 + down_input_scale = torch.empty( + ( + gateup_output.shape[0], + gateup_output.shape[1], + gateup_output.shape[2] // 2 // scale_block_size, + ), + device=gateup_output.device, + dtype=torch.float32, + ) + silu_and_mul_masked_post_quant_fwd( + gateup_output, + down_input, + down_input_scale, + scale_block_size, + masked_m, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + del gateup_output + + # GroupGemm-1 + n = self.w2_weight.size(1) + down_input_fp8 = ( + down_input, + ( + down_input_scale + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + else deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(down_input_scale) + ), + ) + down_output = torch.empty( + (num_groups, m, n), device=down_input.device, dtype=torch.bfloat16 + ) + deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked( + down_input_fp8, + self.w2_weight_fp8, + down_output, + masked_m, + expected_m, + ) + + return down_output + + def forward_npu( + self, + dispatch_output: Union[DeepEPNormalOutput, DeepEPLLOutput], + ): + assert self.quant_method is not None + assert self.moe_runner_config.activation == "silu" + + import torch_npu + + from sglang.srt.layers.moe.token_dispatcher import DispatchOutputChecker + + # NOTE: Ascend's Dispatch & Combine does not support FP16 + output_dtype = torch.bfloat16 + group_list_type = 1 + + def _forward_normal(dispatch_output: DeepEPNormalOutput): + if TYPE_CHECKING: + assert isinstance(dispatch_output, DeepEPNormalOutput) + hidden_states, _, _, num_recv_tokens_per_expert = dispatch_output + + if isinstance(hidden_states, tuple): + per_token_scale = hidden_states[1] + hidden_states = hidden_states[0] + else: + # dynamic quant + hidden_states, per_token_scale = torch_npu.npu_dynamic_quant( + hidden_states + ) + + group_list = torch.tensor(num_recv_tokens_per_expert, dtype=torch.int64).to( + hidden_states.device + ) + + # gmm1: gate_up_proj + hidden_states = torch_npu.npu_grouped_matmul( + x=[hidden_states], + weight=[self.w13_weight], + scale=[self.w13_weight_scale.to(output_dtype)], + per_token_scale=[per_token_scale], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=output_dtype, + )[0] + + # act_fn: swiglu + hidden_states = torch_npu.npu_swiglu(hidden_states) + hidden_states, swiglu_out_scale = torch_npu.npu_dynamic_quant(hidden_states) + + # gmm2: down_proj + hidden_states = torch_npu.npu_grouped_matmul( + x=[hidden_states], + weight=[self.w2_weight], + scale=[self.w2_weight_scale.to(output_dtype)], + per_token_scale=[swiglu_out_scale], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=output_dtype, + )[0] + + return hidden_states + + def _forward_ll(dispatch_output: DeepEPLLOutput): + if TYPE_CHECKING: + assert isinstance(dispatch_output, DeepEPLLOutput) + hidden_states, topk_idx, topk_weights, group_list, _ = dispatch_output + + per_token_scale = hidden_states[1] + hidden_states = hidden_states[0] + + group_list = group_list.to(torch.int64) + + # gmm1: gate_up_proj + hidden_states = torch_npu.npu_grouped_matmul( + x=[hidden_states], + weight=[self.w13_weight], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=torch.int32, + )[0] + + # act_fn: swiglu + hidden_states, swiglu_out_scale = torch_npu.npu_dequant_swiglu_quant( + x=hidden_states, + weight_scale=self.w13_weight_scale.to(torch.float32), + activation_scale=per_token_scale, + bias=None, + quant_scale=None, + quant_offset=None, + group_index=group_list, + activate_left=True, + quant_mode=1, + ) + + # gmm2: down_proj + hidden_states = torch_npu.npu_grouped_matmul( + x=[hidden_states], + weight=[self.w2_weight], + scale=[self.w2_weight_scale.to(output_dtype)], + per_token_scale=[swiglu_out_scale], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=output_dtype, + )[0] + + return hidden_states + + if DispatchOutputChecker.format_is_deepep_normal(dispatch_output): + return _forward_normal(dispatch_output) + elif DispatchOutputChecker.format_is_deepep_ll(dispatch_output): + return _forward_ll(dispatch_output) + else: + raise ValueError(f"Not Supported DeepEP format {dispatch_output.format}") + + +def get_moe_impl_class(quant_config: Optional[QuantizationConfig] = None): + if get_moe_a2a_backend().is_deepep(): + return DeepEPMoE + + # NEW: Direct FP4 detection (bypasses EP requirements) + # Check for FP4 quantization with TRTLLM flag, regardless of EP + if get_moe_runner_backend().is_flashinfer_trtllm(): + # FlashInferFP4MoE must be paired with ModelOptNvFp4FusedMoEMethod. + # If UnquantizedFusedMoEMethod is detected, fall back to FusedMoE instead. + if quant_config is None: + return FusedMoE + try: + # Check the quantization argument directly + quantization = global_server_args_dict.get("quantization") + if quantization == "modelopt_fp4": + from sglang.srt.layers.moe.fused_moe_triton.layer import ( + FlashInferFP4MoE, + ) + + return FlashInferFP4MoE + except: + pass + + if should_use_flashinfer_trtllm_moe(): + return FlashInferFusedMoE + if get_moe_runner_backend().is_flashinfer_cutlass(): + return FusedMoE + if get_moe_expert_parallel_world_size() > 1: + return EPMoE + return FusedMoE diff --git a/python/sglang/srt/layers/moe/fused_moe_native.py b/python/sglang/srt/layers/moe/fused_moe_native.py new file mode 100644 index 000000000..a3d3a09bf --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_native.py @@ -0,0 +1,101 @@ +""" +Torch-native implementation for FusedMoE. This is used for torch.compile. +It is based on https://github.com/pytorch-labs/gpt-fast/blob/32971d3129541c5bfb4f715abc33d1c5f408d204/mixtral-moe/model.py#L204 +""" + +import torch +from torch.nn import functional as F + +from sglang.srt.layers.activation import GeluAndMul, SiluAndMul +from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig +from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput +from sglang.srt.layers.moe.topk import StandardTopKOutput + + +def fused_moe_forward_native( + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, +) -> torch.Tensor: + + x, topk_output = dispatch_output + moe_runner_config = layer.moe_runner_config + + if moe_runner_config.apply_router_weight_on_input: + raise NotImplementedError() + + topk_weights, topk_ids, _ = topk_output + + w13_weights = layer.w13_weight[topk_ids] + w1_weights, w3_weights = torch.chunk(w13_weights, 2, dim=2) + w2_weights = layer.w2_weight[topk_ids] + x1 = torch.einsum("ti,taoi -> tao", x, w1_weights) + if moe_runner_config.activation == "silu": + x1 = F.silu(x1) + elif moe_runner_config.activation == "gelu": + x1 = F.gelu(x1) + else: + raise ValueError(f"Unsupported activation: {moe_runner_config.activation=}") + x3 = torch.einsum("ti, taoi -> tao", x, w3_weights) + expert_outs = torch.einsum("tao, taio -> tai", (x1 * x3), w2_weights) + return torch.einsum("tai,ta -> ti", expert_outs, topk_weights.to(expert_outs.dtype)) + + +def moe_forward_native( + layer: torch.nn.Module, + x: torch.Tensor, + topk_output: StandardTopKOutput, + moe_runner_config: MoeRunnerConfig, +) -> torch.Tensor: + + if moe_runner_config.apply_router_weight_on_input: + raise NotImplementedError() + + topk_weights, topk_ids, _ = topk_output + + # Ref code from https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/e0828e3cc0a03408724b80c3cc92c8e072db8d01/modeling_deepseek.py#L589 + len_experts = layer.num_experts + + cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts)) + cnts.scatter_(1, topk_ids.to(torch.int64), 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + + sorted_tokens = x[idxs // topk_ids.shape[1]] + tokens_per_expert = tokens_per_expert.cpu().numpy() + + if moe_runner_config.activation == "silu": + act = SiluAndMul() + elif moe_runner_config.activation == "gelu": + act = GeluAndMul() + else: + raise ValueError(f"Unsupported activation: {moe_runner_config.activation=}") + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + + layer_w13_weight = layer.w13_weight[i] + layer_w2_weight = layer.w2_weight[i] + + gate_up = F.linear(tokens_for_this_expert, layer_w13_weight) + gate_up = act(gate_up) + expert_out = F.linear(gate_up, layer_w2_weight) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + new_x = torch.empty_like(outs) + + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weights.dtype) + .mul_(topk_weights.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/__init__.py b/python/sglang/srt/layers/moe/fused_moe_triton/__init__.py new file mode 100644 index 000000000..be3ed3af4 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/__init__.py @@ -0,0 +1,42 @@ +from contextlib import contextmanager +from typing import Any, Dict, Optional + +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_experts +from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config import ( + get_config_file_name, + try_get_optimal_moe_config, +) +from sglang.srt.layers.moe.fused_moe_triton.layer import ( + FusedMoE, + FusedMoeWeightScaleSupported, +) +from sglang.srt.layers.moe.fused_moe_triton.moe_align_block_size import ( + moe_align_block_size, +) + +_config: Optional[Dict[str, Any]] = None + + +@contextmanager +def override_config(config): + global _config + old_config = _config + _config = config + yield + _config = old_config + + +def get_config() -> Optional[Dict[str, Any]]: + return _config + + +__all__ = [ + "FusedMoE", + "FusedMoeWeightScaleSupported", + "override_config", + "get_config", + "fused_experts", + "get_config_file_name", + "moe_align_block_size", + "try_get_optimal_moe_config", +] diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/README.md b/python/sglang/srt/layers/moe/fused_moe_triton/configs/README.md new file mode 100644 index 000000000..3679e698a --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/README.md @@ -0,0 +1,15 @@ +This directory contains tuned configurations for different settings of the fused_moe kernel. +For different settings of +- E (number of experts) +- N (intermediate size) +- device_name (torch.cuda.get_device_name()) +- dtype: The data type used by the fused MoE kernel for computation. Supported types include fp8_w8a8, int8_w8a8, int8_w8a16, int4_w4a16, etc. This determines the precision and quantization scheme for both weights and activations. +- block_shape: The block quantization shape introduced starting from DeepSeek V3/R1 models. This parameter defines the granularity for block-wise quantization, typically specified as `[block_n, block_k]` where `block_n` and `block_k` represent the block dimensions. For example, DeepSeek V3 commonly uses `[128, 128]` block shapes for efficient block-wise FP8 quantization. + +the JSON file contains a mapping from M (batch size) to the chosen configuration. + +The example configurations provided are for the Mixtral model for TP2 on H100 +and TP4 on A100. Mixtral has intermediate size N = 14336, i.e. for TP2 we have +N = 7168 and for TP4 we have N = 3584. + +See `benchmark/kernels/fused_moe_triton/README.md` on how to generate these config files. diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=14336,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=14336,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..56c1a4e3a --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=14336,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=14336,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=14336,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..21924c7e7 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=14336,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=1792,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=1792,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..598993c61 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=1792,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "5120": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "9216": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "13312": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "17408": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "25600": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "33792": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "41984": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "50176": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "58368": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..a7c626e9b --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "5120": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "9216": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "17408": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "25600": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "33792": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "41984": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "50176": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "58368": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3072,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3072,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..14bf13935 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3072,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "5120": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "9216": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "17408": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "25600": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "33792": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "41984": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "50176": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "58368": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3072,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3072,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json new file mode 100644 index 000000000..29972c5f3 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3072,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "5120": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "9216": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "17408": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "25600": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "33792": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "41984": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "50176": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "58368": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3072,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3072,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..082fba99b --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3072,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "5120": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "9216": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "17408": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "25600": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "33792": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "41984": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "50176": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "58368": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3584,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3584,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..dbed86769 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3584,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "5120": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "9216": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "13312": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "17408": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "25600": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "33792": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "41984": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "50176": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "58368": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..b89336a26 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "5120": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "9216": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "17408": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "25600": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "33792": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "41984": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "50176": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "58368": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..39bc5b9e5 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "5120": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "9216": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "17408": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "25600": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "33792": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "41984": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "50176": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "58368": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..35d67f837 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=1,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "5120": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "9216": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "17408": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "25600": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "33792": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "41984": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "50176": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "58368": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=144,N=512,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=144,N=512,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..587fb2f2e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=144,N=512,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1024,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1024,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..089894816 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1024,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1024,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1024,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..9814a3819 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1024,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1344,device_name=NVIDIA_A100-SXM4-40GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1344,device_name=NVIDIA_A100-SXM4-40GB.json new file mode 100644 index 000000000..9262a74a4 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1344,device_name=NVIDIA_A100-SXM4-40GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1344,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1344,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..d251f9b5a --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1344,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1344,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1344,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..0ecf814a2 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1344,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=14336,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=14336,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..6443c5cd0 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=14336,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=14336,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=14336,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..1fe9aad96 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=14336,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1792,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1792,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..9b48fe26d --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1792,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "5120": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "9216": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "17408": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "25600": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "33792": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "41984": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "50176": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "58368": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..5ff4a1600 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "5120": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "9216": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "17408": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "25600": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "33792": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "41984": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "50176": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "58368": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..7e09be1ff --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=2688,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=2688,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..039a10ed1 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=2688,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..3793fcafe --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3072,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3072,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..f1f626667 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3072,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3072,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3072,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json new file mode 100644 index 000000000..6c809a096 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3072,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "1536": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "4096": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3200,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3200,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..3a9e2ef0f --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3200,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,130 @@ +{ + "3328": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "768": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "1792": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2560": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "2816": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3584": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "3840": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1280": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "2304": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3584,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3584,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..1bfd02a41 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3584,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..0bb90b7e4 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,218 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "5120": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "9216": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "13312": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "17408": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "25600": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "33792": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "41984": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "50176": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "58368": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=6400,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=6400,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..033ba506d --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=6400,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,130 @@ +{ + "3840": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "1792": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "3584": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "2816": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1280": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "768": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "3328": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "2560": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "2304": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json new file mode 100644 index 000000000..14723c576 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..9e9d27b1c --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json new file mode 100644 index 000000000..6ac5ba746 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "1536": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=800,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=800,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..c87ca0de8 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=16,N=800,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,130 @@ +{ + "2048": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "1792": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "3328": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2560": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 4 + }, + "768": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "2816": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2304": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2 + }, + "1280": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3840": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3584": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=160,N=192,device_name=NVIDIA_A800-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=160,N=192,device_name=NVIDIA_A800-SXM4-80GB.json new file mode 100644 index 000000000..09d812937 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=160,N=192,device_name=NVIDIA_A800-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=20,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=20,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..d471f9821 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=20,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=24,N=1024,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=24,N=1024,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..fe8dab46e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=24,N=1024,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..1cd253d11 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8.json new file mode 100644 index 000000000..9acb73b02 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..ddd55abe0 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json new file mode 100644 index 000000000..4ce62139b --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..2e692a158 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_H20,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_H20,block_shape=[128, 128].json new file mode 100644 index 000000000..4e7a2a1a0 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_H20,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..857d11e48 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..90facc4a8 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=128,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..e25d0492d --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..e25d0492d --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5f32e830a --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..2840e9f47 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c6c98c28e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6fcf40875 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=64,device_name=NVIDIA_A800-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=64,device_name=NVIDIA_A800-SXM4-80GB.json new file mode 100644 index 000000000..4dcfefa5b --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=64,device_name=NVIDIA_A800-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=64,device_name=NVIDIA_L20,dtype=int8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=64,device_name=NVIDIA_L20,dtype=int8_w8a8.json new file mode 100644 index 000000000..ac9349fef --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=64,device_name=NVIDIA_L20,dtype=int8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=64,device_name=NVIDIA_L40S,dtype=int8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=64,device_name=NVIDIA_L40S,dtype=int8_w8a8.json new file mode 100644 index 000000000..b99349fb8 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=256,N=64,device_name=NVIDIA_L40S,dtype=int8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..e9c5b9251 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..8cc6c643f --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_A800-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_A800-SXM4-80GB.json new file mode 100644 index 000000000..283ffd8ff --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_A800-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..e2f8164cc --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..ef6a0479c --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..67bf2b720 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..da71451a7 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=1280,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=2560,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=2560,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..14c6f8c1a --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=2560,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=2560,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=2560,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..2b974a78d --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=2560,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=2560,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=2560,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..054873f38 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=2560,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..869f75127 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..40d85e43e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..4ebd7a781 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..90416662e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=320,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..3483f05a2 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..b2799ed3a --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_A800-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_A800-SXM4-80GB.json new file mode 100644 index 000000000..8a18afe7d --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_A800-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json new file mode 100644 index 000000000..a1ec0b3ce --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..96437b5ba --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..6e6050694 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..824009a59 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..bf2aa0d5f --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=64,N=640,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=AMD_Instinct_MI300X.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=AMD_Instinct_MI300X.json new file mode 100644 index 000000000..4d4b752fa --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=AMD_Instinct_MI300X.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 1, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 1, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 32, + "kpack": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=AMD_Instinct_MI325X.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=AMD_Instinct_MI325X.json new file mode 100644 index 000000000..4d4b752fa --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=AMD_Instinct_MI325X.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 1, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 1, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 32, + "kpack": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=AMD_Radeon_Graphics.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=AMD_Radeon_Graphics.json new file mode 100644 index 000000000..4d4b752fa --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=AMD_Radeon_Graphics.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 1, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 1, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 32, + "kpack": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..3f3ccdafa --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,138 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..62fd21136 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..da40db132 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=14336,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=AMD_Instinct_MI300X.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=AMD_Instinct_MI300X.json new file mode 100644 index 000000000..a218fc406 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=AMD_Instinct_MI300X.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=AMD_Instinct_MI325X.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=AMD_Instinct_MI325X.json new file mode 100644 index 000000000..a218fc406 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=AMD_Instinct_MI325X.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=AMD_Radeon_Graphics.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=AMD_Radeon_Graphics.json new file mode 100644 index 000000000..a218fc406 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=AMD_Radeon_Graphics.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_A100-SXM4-40GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_A100-SXM4-40GB.json new file mode 100644 index 000000000..f4c0f8417 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_A100-SXM4-40GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..5c8185cfd --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..97c9f4445 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..9d0db1cdc --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..1c906494a --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=1792,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..0bb423b28 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..555718733 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..26bcbf269 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..430f50090 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..c51565f67 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=2048,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=AMD_Instinct_MI300X.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=AMD_Instinct_MI300X.json new file mode 100644 index 000000000..3682cc548 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=AMD_Instinct_MI300X.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 32, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=AMD_Instinct_MI325X.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=AMD_Instinct_MI325X.json new file mode 100644 index 000000000..3682cc548 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=AMD_Instinct_MI325X.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 32, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=AMD_Radeon_Graphics.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=AMD_Radeon_Graphics.json new file mode 100644 index 000000000..3682cc548 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=AMD_Radeon_Graphics.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 32, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_A100-SXM4-40GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_A100-SXM4-40GB.json new file mode 100644 index 000000000..b41f9d443 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_A100-SXM4-40GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..edf2a38d1 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json new file mode 100644 index 000000000..7532826a6 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..673bae2ba --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..b2100cebb --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..449a13384 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..d317994cb --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_L40S.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_L40S.json new file mode 100644 index 000000000..6499c8586 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=3584,device_name=NVIDIA_L40S.json @@ -0,0 +1,173 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_ctas": 1, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_ctas": 1, + "num_stages": 7 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 128, + "num_warps": 2, + "num_ctas": 1, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_ctas": 1, + "num_stages": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_ctas": 1, + "num_stages": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_ctas": 1, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_ctas": 1, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_ctas": 1, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_ctas": 1, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_ctas": 1, + "num_stages": 2 + }, + "192": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_ctas": 1, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_ctas": 1, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 128, + "num_warps": 2, + "num_ctas": 1, + "num_stages": 8 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_ctas": 1, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_ctas": 1, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_ctas": 1, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_ctas": 1, + "num_stages": 2 + }, + "6144": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_ctas": 1, + "num_stages": 2 + }, + "8192": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_ctas": 1, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8.json new file mode 100644 index 000000000..d7f14d665 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8.json @@ -0,0 +1,178 @@ +{ + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8192": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16384": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32768": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "65536": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "131072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8.json new file mode 100644 index 000000000..d7f14d665 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8.json @@ -0,0 +1,178 @@ +{ + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8192": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16384": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32768": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "65536": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "131072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8.json new file mode 100644 index 000000000..d7f14d665 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8.json @@ -0,0 +1,178 @@ +{ + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8192": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16384": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32768": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "65536": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "131072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..dbc624731 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..cc614e635 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..32c0c9da4 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..72c3f560b --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..dd07b3f6e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=4096,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=AMD_Instinct_MI300X.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=AMD_Instinct_MI300X.json new file mode 100644 index 000000000..21742854c --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=AMD_Instinct_MI300X.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 32, + "kpack": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=AMD_Instinct_MI325X.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=AMD_Instinct_MI325X.json new file mode 100644 index 000000000..21742854c --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=AMD_Instinct_MI325X.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 32, + "kpack": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=AMD_Radeon_Graphics.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=AMD_Radeon_Graphics.json new file mode 100644 index 000000000..21742854c --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=AMD_Radeon_Graphics.json @@ -0,0 +1,200 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 2, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 32, + "kpack": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..f578c8d01 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..918f68396 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..e341a6791 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..a841518ca --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..13cc2cee1 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=7168,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8.json new file mode 100644 index 000000000..d9d2f5eac --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8.json @@ -0,0 +1,175 @@ +{ + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8192": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16384": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32768": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "65536": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "131072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8.json new file mode 100644 index 000000000..d9d2f5eac --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8.json @@ -0,0 +1,175 @@ +{ + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8192": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16384": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32768": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "65536": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "131072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8.json new file mode 100644 index 000000000..d9d2f5eac --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8.json @@ -0,0 +1,175 @@ +{ + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + }, + "8192": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "16384": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "32768": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "65536": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 1 + }, + "131072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + "kpack": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json new file mode 100644 index 000000000..34b916e57 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..b50cfc13d --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_1_0/E=8,N=8192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_A800-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_A800-SXM4-80GB.json new file mode 100644 index 000000000..952b750aa --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_A800-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..5de5605d4 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_H20.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_H20.json new file mode 100644 index 000000000..2221e99cd --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_H20.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..74374c573 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=192,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..663f32539 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..b34b6e4e8 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H20.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H20.json new file mode 100644 index 000000000..ab169a018 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H20.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..b11acb1ee --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..ab6e15552 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=384,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=512,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=512,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..249359fb9 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=512,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_A800-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_A800-SXM4-80GB.json new file mode 100644 index 000000000..4d3fcf75c --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_A800-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..9259aa780 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..b4efc9b7e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H20.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H20.json new file mode 100644 index 000000000..03dfc73b6 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H20.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..e7f9143f5 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..beaac7f64 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=768,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=96,device_name=NVIDIA_H20.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=96,device_name=NVIDIA_H20.json new file mode 100644 index 000000000..ebff99e26 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=128,N=96,device_name=NVIDIA_H20.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=129,N=352,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=129,N=352,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json new file mode 100644 index 000000000..ff3fc2d01 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=129,N=352,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=160,N=320,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=160,N=320,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json new file mode 100644 index 000000000..550d314f0 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=160,N=320,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=161,N=192,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=161,N=192,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json new file mode 100644 index 000000000..763a1657f --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=161,N=192,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..63ba315b0 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..42d00fe60 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..739979fad --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..0762b9084 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..29b5aab37 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json new file mode 100644 index 000000000..eb5b9b741 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..741060e32 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..f3f8f7198 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..52b459d2e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=264,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json new file mode 100644 index 000000000..da6cf2a04 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..e078fd083 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5ed7ad20f --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=64,device_name=NVIDIA_A800-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=64,device_name=NVIDIA_A800-SXM4-80GB.json new file mode 100644 index 000000000..b00455054 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=272,N=64,device_name=NVIDIA_A800-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=288,N=64,device_name=NVIDIA_A800-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=288,N=64,device_name=NVIDIA_A800-SXM4-80GB.json new file mode 100644 index 000000000..f8f336e26 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=288,N=64,device_name=NVIDIA_A800-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..e341a6791 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_2_0/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_0/E=16,N=1024,device_name=NVIDIA_B200.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_0/E=16,N=1024,device_name=NVIDIA_B200.json new file mode 100644 index 000000000..1fa444bca --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_0/E=16,N=1024,device_name=NVIDIA_B200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=352,device_name=NVIDIA_RTX_6000_Ada_Generation,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=352,device_name=NVIDIA_RTX_6000_Ada_Generation,dtype=fp8_w8a8.json new file mode 100644 index 000000000..f8fd97b5e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=352,device_name=NVIDIA_RTX_6000_Ada_Generation,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c4751c00f --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6d6af0808 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=768,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=768,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..379708af4 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=768,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=768,device_name=NVIDIA_H20.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=768,device_name=NVIDIA_H20.json new file mode 100644 index 000000000..1f3648ca8 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=128,N=768,device_name=NVIDIA_H20.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..adeaacb0e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=320,device_name=NVIDIA_H20-3e.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=320,device_name=NVIDIA_H20-3e.json new file mode 100644 index 000000000..a3022a054 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=320,device_name=NVIDIA_H20-3e.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=384,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=384,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..c3c6e0ac0 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=384,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..1eba1e4b8 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=160,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..8e9668889 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..3bac45f89 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c9c8bb66e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..68f6fb5ac --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..8ddf8beae --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..532c16e89 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=384,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=384,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5ee3d4d35 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=384,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=384,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=384,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..358873315 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=384,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=384,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=384,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..02a119508 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=384,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=385,N=128,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=385,N=128,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..faf1aa4d4 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=385,N=128,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=385,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=385,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..b5c45dd72 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=385,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..6f586d1fa --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_3_1/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=384,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=384,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c848ea577 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=384,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=129,N=352,device_name=NVIDIA_B200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=129,N=352,device_name=NVIDIA_B200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..41d97b17b --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=129,N=352,device_name=NVIDIA_B200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "4096": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=129,N=352,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=129,N=352,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition,dtype=fp8_w8a8.json new file mode 100644 index 000000000..f8fd97b5e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=129,N=352,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=129,N=704,device_name=NVIDIA_B200,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=129,N=704,device_name=NVIDIA_B200,dtype=fp8_w8a8.json new file mode 100644 index 000000000..b962d1950 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=129,N=704,device_name=NVIDIA_B200,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=161,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition,dtype=fp8_w8a8.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=161,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition,dtype=fp8_w8a8.json new file mode 100644 index 000000000..f8fd97b5e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=161,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition,dtype=fp8_w8a8.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=256,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=256,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..4e36c1544 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=256,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=257,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=257,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a6c635be4 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=257,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=257,N=64,device_name=NVIDIA_A100-SXM4-80GB.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=257,N=64,device_name=NVIDIA_A100-SXM4-80GB.json new file mode 100644 index 000000000..dc8d6d68b --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=257,N=64,device_name=NVIDIA_A100-SXM4-80GB.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=384,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=384,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..486b94720 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=384,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..b8f35b62e --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H20-3e.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H20-3e.json new file mode 100644 index 000000000..039d5ade7 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H20-3e.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=256,device_name=NVIDIA_H20-3e.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=256,device_name=NVIDIA_H20-3e.json new file mode 100644 index 000000000..991b315f7 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=256,device_name=NVIDIA_H20-3e.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=64,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=64,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..64861b390 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=64,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py new file mode 100644 index 000000000..3a64c1457 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py @@ -0,0 +1,701 @@ +# NOTE: this file will be separated into sglang/srt/layers/moe/moe_runner/triton_utils.py +# Adapted from https://github.com/vllm-project/vllm/blob/a6221a144af772fd1a68fe7e627935dc53e81738/vllm/model_executor/layers/fused_moe/fused_moe.py + +"""Fused MoE kernel.""" + +from __future__ import annotations + +import functools +import os +from typing import TYPE_CHECKING, List, Optional + +import torch +import triton.language as tl + +from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig +from sglang.srt.utils import ( + cpu_has_amx_support, + direct_register_custom_op, + get_bool_env_var, + is_cpu, + is_cuda, + is_hip, +) + +from .fused_moe_triton_config import get_config_dtype_str, try_get_optimal_moe_config +from .fused_moe_triton_kernels import invoke_fused_moe_kernel, moe_sum_reduce_triton +from .moe_align_block_size import moe_align_block_size + +if TYPE_CHECKING: + from sglang.srt.layers.moe.topk import StandardTopKOutput + +_is_hip = is_hip() +_is_cuda = is_cuda() +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _is_cuda: + from sgl_kernel import gelu_and_mul, silu_and_mul +elif _is_cpu and _is_cpu_amx_available: + pass +elif _is_hip: + from sgl_kernel import gelu_and_mul, silu_and_mul + + if _use_aiter: + try: + from aiter import moe_sum + except ImportError: + raise ImportError("aiter is required when SGLANG_USE_AITER is set to True") + else: + from vllm import _custom_ops as vllm_ops + +padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0 + + +def inplace_fused_experts( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + b1: Optional[torch.Tensor] = None, + b2: Optional[torch.Tensor] = None, + activation: int = 0,#0 silu 1 gelu + apply_router_weight_on_input: bool = False, + use_fp8_w8a8: bool = False, + use_int8_w8a8: bool = False, + use_int8_w8a16: bool = False, + use_int4_w4a16: bool = False, + per_channel_quant: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + w1_zp: Optional[torch.Tensor] = None, + w2_zp: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[List[int]] = None, + routed_scaling_factor: Optional[float] = None, + gemm1_alpha: Optional[float] = None, + gemm1_limit: Optional[float] = None, +) -> None: + if isinstance(activation, int): + activation = "silu" if activation == 0 else "gelu" + fused_experts_impl( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + b1, + b2, + True, + activation, + apply_router_weight_on_input, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + use_int4_w4a16, + per_channel_quant, + w1_scale, + w2_scale, + w1_zp, + w2_zp, + a1_scale, + a2_scale, + block_shape, + False, + routed_scaling_factor, + gemm1_alpha, + gemm1_limit, + ) + + +def inplace_fused_experts_fake( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + b1: Optional[torch.Tensor] = None, + b2: Optional[torch.Tensor] = None, + activation: int = 0,#0 silu 1 gelu + apply_router_weight_on_input: bool = False, + use_fp8_w8a8: bool = False, + use_int8_w8a8: bool = False, + use_int8_w8a16: bool = False, + use_int4_w4a16: bool = False, + per_channel_quant: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + w1_zp: Optional[torch.Tensor] = None, + w2_zp: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[List[int]] = None, + routed_scaling_factor: Optional[float] = None, + gemm1_alpha: Optional[float] = None, + gemm1_limit: Optional[float] = None, +) -> None: + pass + + +direct_register_custom_op( + op_name="inplace_fused_experts", + op_func=inplace_fused_experts, + mutates_args=["hidden_states"], + fake_impl=inplace_fused_experts_fake, +) + + +def outplace_fused_experts( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + b1: Optional[torch.Tensor] = None, + b2: Optional[torch.Tensor] = None, + activation: int = 0,#0 silu 1 gelu + apply_router_weight_on_input: bool = False, + use_fp8_w8a8: bool = False, + use_int8_w8a8: bool = False, + use_int8_w8a16: bool = False, + use_int4_w4a16: bool = False, + per_channel_quant: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + w1_zp: Optional[torch.Tensor] = None, + w2_zp: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[List[int]] = None, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + gemm1_alpha: Optional[float] = None, + gemm1_limit: Optional[float] = None, +) -> torch.Tensor: + if isinstance(activation, int): + activation = "silu" if activation == 0 else "gelu" + return fused_experts_impl( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + b1, + b2, + False, + activation, + apply_router_weight_on_input, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + use_int4_w4a16, + per_channel_quant, + w1_scale, + w2_scale, + w1_zp, + w2_zp, + a1_scale, + a2_scale, + block_shape, + no_combine=no_combine, + routed_scaling_factor=routed_scaling_factor, + gemm1_alpha=gemm1_alpha, + gemm1_limit=gemm1_limit, + ) + + +def outplace_fused_experts_fake( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + b1: Optional[torch.Tensor] = None, + b2: Optional[torch.Tensor] = None, + activation: int = 0,#0 silu 1 gelu + apply_router_weight_on_input: bool = False, + use_fp8_w8a8: bool = False, + use_int8_w8a8: bool = False, + use_int8_w8a16: bool = False, + use_int4_w4a16: bool = False, + per_channel_quant: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + w1_zp: Optional[torch.Tensor] = None, + w2_zp: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[List[int]] = None, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + gemm1_alpha: Optional[float] = None, + gemm1_limit: Optional[float] = None, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +direct_register_custom_op( + op_name="outplace_fused_experts", + op_func=outplace_fused_experts, + mutates_args=[], + fake_impl=outplace_fused_experts_fake, +) + + +def fused_experts( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_output: StandardTopKOutput, + moe_runner_config: MoeRunnerConfig, + b1: Optional[torch.Tensor] = None, + b2: Optional[torch.Tensor] = None, + use_fp8_w8a8: bool = False, + use_int8_w8a8: bool = False, + use_int8_w8a16: bool = False, + use_int4_w4a16: bool = False, + per_channel_quant: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + w1_zp: Optional[torch.Tensor] = None, + w2_zp: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[List[int]] = None, +): + topk_weights, topk_ids, _ = topk_output + # 统一把 activation -> int（0=silu, 1=gelu） + act_id = ( + 0 if ( + moe_runner_config.activation == 0 + or (isinstance(moe_runner_config.activation, str) + and moe_runner_config.activation.lower() == "silu") + ) else 1 + ) + if moe_runner_config.inplace: + assert not moe_runner_config.no_combine, "no combine + inplace makes no sense" + torch.ops.sglang.inplace_fused_experts( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + b1, + b2, + # moe_runner_config.activation, + act_id, + moe_runner_config.apply_router_weight_on_input, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + use_int4_w4a16, + per_channel_quant, + w1_scale, + w2_scale, + w1_zp, + w2_zp, + a1_scale, + a2_scale, + block_shape, + moe_runner_config.routed_scaling_factor, + moe_runner_config.gemm1_alpha, + moe_runner_config.gemm1_clamp_limit, + ) + return hidden_states + else: + return torch.ops.sglang.outplace_fused_experts( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + b1, + b2, + # moe_runner_config.activation, + act_id, + moe_runner_config.apply_router_weight_on_input, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + use_int4_w4a16, + per_channel_quant, + w1_scale, + w2_scale, + w1_zp, + w2_zp, + a1_scale, + a2_scale, + block_shape, + no_combine=moe_runner_config.no_combine, + routed_scaling_factor=moe_runner_config.routed_scaling_factor, + gemm1_alpha=moe_runner_config.gemm1_alpha, + gemm1_limit=moe_runner_config.gemm1_clamp_limit, + ) + + +@torch.compile +def moe_sum_reduce_torch_compile(x, out, routed_scaling_factor): + torch.sum(x, dim=1, out=out) + out.mul_(routed_scaling_factor) + + +@torch.compile +def swiglu_with_alpha_and_limit(x, gemm1_alpha, gemm1_limit): + gate, up = x[..., ::2], x[..., 1::2] + gate = gate.clamp(min=None, max=gemm1_limit) + up = up.clamp(min=-gemm1_limit, max=gemm1_limit) + return gate * torch.sigmoid(gate * gemm1_alpha) * (up + 1) + + +def fused_experts_impl( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + b1: Optional[torch.Tensor] = None, + b2: Optional[torch.Tensor] = None, + inplace: bool = False, + activation: int = 0,#0 silu 1 gelu + apply_router_weight_on_input: bool = False, + use_fp8_w8a8: bool = False, + use_int8_w8a8: bool = False, + use_int8_w8a16: bool = False, + use_int4_w4a16: bool = False, + per_channel_quant: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + w1_zp: Optional[torch.Tensor] = None, + w2_zp: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[List[int]] = None, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + gemm1_alpha: Optional[float] = None, + gemm1_limit: Optional[float] = None, +): + padded_size = padding_size + if not (use_fp8_w8a8 or use_int8_w8a8) or block_shape is not None or _use_aiter: + padded_size = 0 + + # Check constraints. + if use_int4_w4a16: + assert hidden_states.shape[1] // 2 == w1.shape[2], "Hidden size mismatch" + else: + assert ( + hidden_states.shape[1] == w1.shape[2] - padded_size + ), f"Hidden size mismatch" + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" + assert w1.is_contiguous(), "Expert weights1 must be contiguous" + assert w2.is_contiguous(), "Expert weights2 must be contiguous" + assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16] + + num_tokens, _ = hidden_states.shape + E, N, _ = w1.shape + # We execute the fused_moe kernel in chunks to circumvent this issue: + # https://github.com/vllm-project/vllm/issues/5938 + CHUNK_SIZE = 64 * 1024 + M = min(num_tokens, CHUNK_SIZE) + config_dtype = get_config_dtype_str( + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + use_int4_w4a16=use_int4_w4a16, + dtype=hidden_states.dtype, + ) + + get_config_func = functools.partial( + try_get_optimal_moe_config, + w1.shape, + (w2.shape[0], w2.shape[1], w2.shape[2] - padded_size), + topk_ids.shape[1], + config_dtype, + block_shape=block_shape, + ) + + config = get_config_func(M) + + cache = torch.empty( + M * topk_ids.shape[1] * max(N, w2.shape[1]), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + intermediate_cache1 = cache[: M * topk_ids.shape[1] * N].view( + (M, topk_ids.shape[1], N), + ) + intermediate_cache2 = torch.empty( + (M * topk_ids.shape[1], N // 2), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + intermediate_cache3 = cache[: M * topk_ids.shape[1] * w2.shape[1]].view( + (M, topk_ids.shape[1], w2.shape[1]), + ) + + compute_type = tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16 + + if no_combine: + assert not inplace + out_hidden_states = torch.empty( + (num_tokens, topk_ids.shape[1], w2.shape[1]), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + elif inplace: + out_hidden_states = hidden_states + else: + out_hidden_states = torch.empty_like(hidden_states) + + for chunk in range((num_tokens // CHUNK_SIZE) + 1): + begin_chunk_idx, end_chunk_idx = ( + chunk * CHUNK_SIZE, + min((chunk + 1) * CHUNK_SIZE, num_tokens), + ) + curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx] + tokens_in_chunk, _ = curr_hidden_states.shape + + if tokens_in_chunk == 0: + break + + if tokens_in_chunk < CHUNK_SIZE and chunk > 0: + # Adjust the intermediate cache size and config for the last + # chunk. Note that in most cases we only have one chunk + # so the cache size and config are already set correctly and + # do not need to be adjusted. + intermediate_cache1 = intermediate_cache1[:tokens_in_chunk] + intermediate_cache2 = intermediate_cache2[ + : tokens_in_chunk * topk_ids.shape[1] + ] + intermediate_cache3 = intermediate_cache3[:tokens_in_chunk] + config = get_config_func(tokens_in_chunk) + + curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx] + curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx] + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + curr_topk_ids, config["BLOCK_SIZE_M"], E + ) + + invoke_fused_moe_kernel( + curr_hidden_states, + w1, + b1, + intermediate_cache1, + a1_scale, + w1_scale, + w1_zp, + curr_topk_weights, + curr_topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + apply_router_weight_on_input, + topk_ids.shape[1], + config, + compute_type=compute_type, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + use_int4_w4a16=use_int4_w4a16, + per_channel_quant=per_channel_quant, + block_shape=block_shape, + ) + if activation == "silu": + if gemm1_alpha is not None: + assert gemm1_limit is not None + intermediate_cache2 = swiglu_with_alpha_and_limit( + intermediate_cache1.view(-1, N), + gemm1_alpha, + gemm1_limit, + ) + elif _is_cuda or _is_hip: + silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2) + else: + vllm_ops.silu_and_mul( + intermediate_cache2, intermediate_cache1.view(-1, N) + ) + elif activation == "gelu": + assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu" + assert gemm1_limit is None, "gemm1_limit is not supported for gelu" + if _is_cuda or _is_hip: + gelu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2) + else: + vllm_ops.gelu_and_mul( + intermediate_cache2, intermediate_cache1.view(-1, N) + ) + else: + raise ValueError(f"Unsupported activation: {activation=}") + + invoke_fused_moe_kernel( + intermediate_cache2, + w2, + b2, + ( + intermediate_cache3 + if not no_combine and topk_ids.shape[1] != 1 + else out_hidden_states[begin_chunk_idx:end_chunk_idx].unsqueeze(0) + ), + a2_scale, + w2_scale, + w2_zp, + curr_topk_weights, + curr_topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + not apply_router_weight_on_input, + 1, + config, + compute_type=compute_type, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + use_int4_w4a16=use_int4_w4a16, + per_channel_quant=per_channel_quant, + block_shape=block_shape, + ) + + if routed_scaling_factor is None: + routed_scaling_factor = 1.0 + + if no_combine: + pass + elif _is_cuda: + if topk_ids.shape[1] == 1 and routed_scaling_factor == 1.0: + pass # we write directly into out_hidden_states + elif topk_ids.shape[1] == 2 and routed_scaling_factor == 1.0: + torch.add( + intermediate_cache3[:, 0], + intermediate_cache3[:, 1], + out=out_hidden_states[begin_chunk_idx:end_chunk_idx], + ).squeeze(dim=1) + else: + # According to micro benchmark results, torch.compile can get better performance for small token. + if tokens_in_chunk <= 32: + moe_sum_reduce_torch_compile( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states[begin_chunk_idx:end_chunk_idx], + routed_scaling_factor, + ) + else: + moe_sum_reduce_triton( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states[begin_chunk_idx:end_chunk_idx], + routed_scaling_factor, + ) + elif _is_hip: + if _use_aiter: + moe_sum( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states[begin_chunk_idx:end_chunk_idx], + ) + else: + # According to micro benchmark results, torch.compile can get better performance for small token. + if tokens_in_chunk <= 32: + moe_sum_reduce_torch_compile( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states[begin_chunk_idx:end_chunk_idx], + routed_scaling_factor, + ) + else: + moe_sum_reduce_triton( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states[begin_chunk_idx:end_chunk_idx], + routed_scaling_factor, + ) + else: + vllm_ops.moe_sum( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states[begin_chunk_idx:end_chunk_idx], + ) + + return out_hidden_states + + +def fused_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_output: StandardTopKOutput, + moe_runner_config: MoeRunnerConfig = MoeRunnerConfig(), + b1: Optional[torch.Tensor] = None, + b2: Optional[torch.Tensor] = None, + use_fp8_w8a8: bool = False, + use_int8_w8a8: bool = False, + use_int8_w8a16: bool = False, + use_int4_w4a16: bool = False, + per_channel_quant: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + w1_zp: Optional[torch.Tensor] = None, + w2_zp: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[List[int]] = None, +) -> torch.Tensor: + """ + This function computes a Mixture of Experts (MoE) layer using two sets of + weights, w1 and w2, and top-k gating mechanism. + + Parameters: + - hidden_states (torch.Tensor): The input tensor to the MoE layer. + - w1 (torch.Tensor): The first set of expert weights. + - w2 (torch.Tensor): The second set of expert weights. + - topk_output (StandardTopKOutput): The top-k output of the experts. + - moe_runner_config (MoeRunnerConfig): The configuration for the MoE runner. + - b1 (Optional[torch.Tensor]): Optional bias for w1. + - b2 (Optional[torch.Tensor]): Optional bias for w2. + - use_fp8_w8a8 (bool): If True, use fp8 arithmetic to compute the inner + products for w1 and w2. Defaults to False. + - use_int8_w8a8 (bool): If True, use int8 arithmetic to compute the inner + products for w1 and w2. Defaults to False. + - use_int8_w8a16 (bool): If True, use fp8 arithmetic to compute the inner + products for w1 and w2. Defaults to False. + - use_int4_w4a16 (bool): If True, use matmul of int4 weight and bf16/fp16 + activation to compute the inner products for w1 and w2. + Defaults to False. + - w1_scale (Optional[torch.Tensor]): Optional scale to be used for + w1. + - w2_scale (Optional[torch.Tensor]): Optional scale to be used for + w2. + - a1_scale (Optional[torch.Tensor]): Optional scale to be used for + a1. + - a2_scale (Optional[torch.Tensor]): Optional scale to be used for + a2. + - block_shape: (Optional[List[int]]): Optional block size for block-wise + quantization. + - gemm1_alpha (Optional[float]): Optional gemm1_alpha for the activation + function. + - gemm1_limit (Optional[float]): Optional gemm1_limit for the swiglu activation + function. + + Returns: + - torch.Tensor: The output tensor after applying the MoE layer. + """ + + return fused_experts( + hidden_states, + w1, + w2, + topk_output, + moe_runner_config=moe_runner_config, + b1=b1, + b2=b2, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + use_int4_w4a16=use_int4_w4a16, + per_channel_quant=per_channel_quant, + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_zp=w1_zp, + w2_zp=w2_zp, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py new file mode 100644 index 000000000..51114aade --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py @@ -0,0 +1,212 @@ +from __future__ import annotations + +import functools +import json +import logging +import os +from typing import Any, Dict, List, Optional, Tuple + +import torch +import triton + +from sglang.srt.utils import get_device_name, is_hip + +logger = logging.getLogger(__name__) +_is_hip = is_hip() + + +def get_config_file_name( + E: int, N: int, dtype: Optional[str], block_shape: Optional[int] = None +) -> str: + device_name = get_device_name().replace(" ", "_") + dtype_selector = "" if not dtype else f",dtype={dtype}" + block_shape_selector = ( + "" if not block_shape or not all(block_shape) else f",block_shape={block_shape}" + ) + return f"E={E},N={N},device_name={device_name}{dtype_selector}{block_shape_selector}.json" + + +@functools.lru_cache +def get_moe_configs( + E: int, + N: int, + dtype: Optional[str], + block_n: Optional[int] = 0, + block_k: Optional[int] = 0, +) -> Optional[Dict[int, Any]]: + """ + Return optimized configurations for the fused MoE kernel. + + The return value will be a dictionary that maps an irregular grid of + batch sizes to configurations of the fused_moe kernel. To evaluate the + kernel on a given batch size bs, the closest batch size in the grid should + be picked and the associated configuration chosen to invoke the kernel. + """ + # Supported Triton versions, should be sorted from the newest to the oldest + supported_triton_versions = ["3.3.1", "3.2.0", "3.1.0"] + + # First look up if an optimized configuration is available in the configs + # directory + json_file_name = get_config_file_name(E, N, dtype, [block_n, block_k]) + + # We found that using the fused_moe_kernel config from Triton 3.1.0 with Triton 3.2.0 results in negative performance gains, + # so we also include the Triton version as a key for finding the fused_moe_kernel config to achieve the best performance. + triton_version = triton.__version__ + version_dir = f"triton_{triton_version.replace('.', '_')}" + config_file_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "configs", + version_dir, + json_file_name, + ) + if os.path.exists(config_file_path): + with open(config_file_path) as f: + # Please note that although we find the config files, performance might still be suboptimal. + # This is because the tuning environment might differ from your current environment. + # For example, updating the Triton version might cause all old configs to become suboptimal. + # To achieve the best performance, consider re-tuning the Triton fused MOE kernel in your environment. + # For the tuning method, refer to: https://github.com/sgl-project/sglang/tree/main/benchmark/kernels/fused_moe_triton + logger.info(f"Using MoE kernel config from {config_file_path}.") + # If a configuration has been found, return it + return {int(key): val for key, val in json.load(f).items()} + + # Searching for other triton versions that supports the same config + for try_triton_version in supported_triton_versions: + if try_triton_version == triton_version: + continue + try_config_file_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "configs", + f"triton_{try_triton_version.replace('.', '_')}", + json_file_name, + ) + if os.path.exists(try_config_file_path): + with open(try_config_file_path) as f: + logger.warning( + f"Config file not found at {config_file_path}. Fallback to triton version {try_triton_version} and use MoE kernel config from {try_config_file_path}. Performance might be sub-optimal!", + ) + # If a configuration has been found, return it + return {int(key): val for key, val in json.load(f).items()} + + # If no optimized configuration is available, we will use the default + # configuration + logger.warning( + ( + "Using default MoE kernel config. Performance might be sub-optimal! " + "Config file not found at %s, you can create them with https://github.com/sgl-project/sglang/tree/main/benchmark/kernels/fused_moe_triton" + ), + config_file_path, + ) + return None + + +def get_default_config( + M: int, + E: int, + N: int, + K: int, + topk: int, + dtype: Optional[str], + is_marlin: bool, + block_shape: Optional[List[int]] = None, +) -> Dict[str, int]: + if dtype == "fp8_w8a8": + if block_shape is None: + config = { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 if _is_hip else 4, + } + if M <= E: + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 if _is_hip else 4, + } + else: + # Block-wise quant: BLOCK_SIZE_K must be divisible by block_shape[1] + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": block_shape[0], + "BLOCK_SIZE_K": block_shape[1], + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 if _is_hip else 3, + } + else: + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + } + # A heuristic: fused marlin works faster with this config for small M + if M <= E or (is_marlin and M <= 32): + config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + } + return config + + +def try_get_optimal_moe_config( + w1_shape: Tuple[int, ...], + w2_shape: Tuple[int, ...], + top_k: int, + dtype: Optional[str], + M: int, + is_marlin: bool = False, + block_shape: Optional[List[int]] = None, +): + from sglang.srt.layers.moe.fused_moe_triton import get_config + + override_config = get_config() + if override_config: + config = override_config + else: + # First try to load optimal config from the file + E, _, N = w2_shape + block_n = block_shape[0] if block_shape else 0 + block_k = block_shape[1] if block_shape else 0 + configs = get_moe_configs(E, N, dtype, block_n, block_k) + + if configs: + # If an optimal configuration map has been found, look up the + # optimal config + config = configs[min(configs.keys(), key=lambda x: abs(x - M))] + else: + # Else use the default config + config = get_default_config( + M, E, N, w1_shape[2], top_k, dtype, is_marlin, block_shape + ) + return config + + +def get_config_dtype_str( + dtype: torch.dtype, + use_int8_w8a16: Optional[bool] = False, + use_int4_w4a16: Optional[bool] = False, + use_fp8_w8a8: Optional[bool] = False, + use_int8_w8a8: Optional[bool] = False, +): + if use_fp8_w8a8: + return "fp8_w8a8" + elif use_int8_w8a8: + return "int8_w8a8" + elif use_int4_w4a16: + return "int4_w4a16" + elif use_int8_w8a16: + return "int8_w8a16" + elif dtype == torch.float: + # avoiding cases where kernel fails when float32 MoE + # use fp16/bfloat16 configs + return "float32" + return None diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py new file mode 100644 index 000000000..6a7229a9b --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py @@ -0,0 +1,799 @@ +from __future__ import annotations + +import os +from typing import Any, Dict, List, Optional + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.quantization.fp8_kernel import ( + per_token_group_quant_fp8, + scaled_fp8_quant, + sglang_per_token_group_quant_fp8, +) +from sglang.srt.layers.quantization.int8_kernel import ( + per_token_group_quant_int8, + per_token_quant_int8, + sglang_per_token_group_quant_int8, +) +from sglang.srt.utils import ( + cpu_has_amx_support, + get_bool_env_var, + is_cpu, + is_cuda, + is_hip, +) + +_is_hip = is_hip() +_is_cuda = is_cuda() +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _is_cuda: + pass +elif _is_cpu and _is_cpu_amx_available: + pass +elif _is_hip: + pass + +padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0 + + +@triton.jit +def write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, +): + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +@triton.jit +def fused_moe_kernel_gptq_awq( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + b_scale_ptr, + b_zp_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + # Matrix dimensions + N: tl.constexpr, + K: tl.constexpr, + EM, + num_valid_tokens, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_am, + stride_ak, + stride_be, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_bse, + stride_bsk, + stride_bsn, + stride_bze, + stride_bzk, + stride_bzn, + group_size: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, + has_zp: tl.constexpr, + use_int4_w4a16: tl.constexpr, + use_int8_w8a16: tl.constexpr, + even_Ks: tl.constexpr, +): + """ + Implements the fused computation for a Mixture of Experts (MOE) using + token and expert matrices. + Key Parameters: + - A: The input tensor representing tokens with shape (*, K), where '*' can + be any shape representing batches and K is the feature dimension of + each token. + - B: The stacked MOE weight tensor with shape (E, N, K), where E is + the number of experts, K is the input feature dimension, and N is + the output feature dimension. + - C: The output cache tensor with shape (M, topk, N), where M is the + total number of tokens post padding, topk is the number of times + each token is repeated, and N is the output feature dimension. + - sorted_token_ids: A tensor containing the sorted indices of tokens, + repeated topk times and arranged by the expert index they are + assigned to. + - expert_ids: A tensor containing the indices of the expert for each + block. It determines which expert matrix from B should be used for + each block in A. + This kernel performs the multiplication of a token by its corresponding + expert matrix as determined by `expert_ids`. The sorting of + `sorted_token_ids` by expert index and padding ensures divisibility by + BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix + multiplication across different blocks processed by the same expert. + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, 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 % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id) + token_mask = offs_token < num_valid_tokens + + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + if off_experts == -1: + # ----------------------------------------------------------- + # Write back zeros to the output when the expert is not + # in the current expert parallel rank. + write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, + ) + return + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + + if use_int4_w4a16: + b_ptrs = ( + b_ptr + + off_experts * stride_be + + (offs_k[:, None] // 2) * stride_bk + + offs_bn[None, :] * stride_bn + ) + b_shifter = (offs_k[:, None] % 2) * 4 + elif use_int8_w8a16: + b_ptrs = ( + b_ptr + + off_experts * stride_be + + offs_k[:, None] * stride_bk + + offs_bn[None, :] * stride_bn + ) + + if not has_zp and use_int4_w4a16: + b_zp_num = 8 + if not has_zp and use_int8_w8a16: + b_zp_num = 128 + elif has_zp and use_int4_w4a16: + b_zp_shifter = (offs_bn[None, :] % 2) * 4 + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the + # K dimension. + + if not even_Ks: + k_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K + k_other = 0.0 + else: + k_mask = None + k_other = None + + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs) + if use_int4_w4a16: + b = (b >> b_shifter) & 0xF + + b_scale_ptrs = ( + b_scale_ptr + + off_experts * stride_bse + + offs_bn[None, :] * stride_bsn + + ((offs_k[:, None] + BLOCK_SIZE_K * k) // group_size) * stride_bsk + ) + b_scale = tl.load(b_scale_ptrs, mask=k_mask, other=k_other) + b_scale = b_scale.to(tl.float32) + + if has_zp and use_int4_w4a16: + offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size + b_zp_ptrs = ( + b_zp_ptr + + off_experts * stride_bze + + (offs_bn[None, :] // 2) * stride_bzn + + offs_k_true * stride_bzk + ) + b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other) + b_zp = (b_zp >> b_zp_shifter) & 0xF + b_zp = b_zp.to(tl.float32) + elif has_zp and use_int8_w8a16: + offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size + b_zp_ptrs = ( + b_zp_ptr + + off_experts * stride_bze + + offs_bn[None, :] * stride_bzn + + offs_k_true * stride_bzk + ) + b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other) + b_zp = b_zp.to(tl.float32) + + # We accumulate along the K dimension. + if has_zp: + b = ((b.to(tl.float32) - b_zp) * b_scale).to(compute_type) + else: + b = ((b.to(tl.float32) - b_zp_num) * b_scale).to(compute_type) + accumulator = tl.dot(a, b, acc=accumulator) + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + if use_int4_w4a16: + b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk + else: + b_ptrs += BLOCK_SIZE_K * stride_bk + + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0) + accumulator = accumulator * moe_weight[:, None] + + accumulator = accumulator.to(compute_type) + # ----------------------------------------------------------- + # Write back the block of the output + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +@triton.jit +def fused_moe_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + bias_ptr, + c_ptr, + a_scale_ptr, + b_scale_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + # Matrix dimensions + N, + K, + EM, + num_valid_tokens, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_am, + stride_ak, + stride_be, + stride_bk, + stride_bn, + stride_bias_e, + stride_bias_n, + stride_cm, + stride_cn, + stride_asm, + stride_ask, + stride_bse, + stride_bsk, + stride_bsn, + # Block size for block-wise quantization + group_n: tl.constexpr, + group_k: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, + use_fp8_w8a8: tl.constexpr, + use_int8_w8a8: tl.constexpr, + use_int8_w8a16: tl.constexpr, + per_channel_quant: tl.constexpr, + even_Ks: tl.constexpr, +): + """ + Implements the fused computation for a Mixture of Experts (MOE) using + token and expert matrices. + + Key Parameters: + - A: The input tensor representing tokens with shape (*, K), where '*' can + be any shape representing batches and K is the feature dimension of + each token. + - B: The stacked MOE weight tensor with shape (E, N, K), where E is + the number of experts, K is the input feature dimension, and N is + the output feature dimension. + - C: The output cache tensor with shape (M, topk, N), where M is the + total number of tokens post padding, topk is the number of times + each token is repeated, and N is the output feature dimension. + - sorted_token_ids: A tensor containing the sorted indices of tokens, + repeated topk times and arranged by the expert index they are + assigned to. + - expert_ids: A tensor containing the indices of the expert for each + block. It determines which expert matrix from B should be used for + each block in A. + + This kernel performs the multiplication of a token by its corresponding + expert matrix as determined by `expert_ids`. The sorting of + `sorted_token_ids` by expert index and padding ensures divisibility by + BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix + multiplication across different blocks processed by the same expert. + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, 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 % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id) + offs_token = offs_token.to(tl.int64) + token_mask = offs_token < num_valid_tokens + + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + + if off_experts == -1: + # ----------------------------------------------------------- + # Write back zeros to the output when the expert is not + # in the current expert parallel rank. + write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, + ) + return + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + + b_ptrs = ( + b_ptr + + off_experts * stride_be + + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + ) + if bias_ptr is not None: + bias = tl.load( + bias_ptr + off_experts * stride_bias_e + offs_bn[None, :] * stride_bias_n + ) + if use_int8_w8a16: + b_scale_ptrs = ( + b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + ) + b_scale = tl.load(b_scale_ptrs) + + if use_fp8_w8a8 or use_int8_w8a8: + # block-wise + if group_k > 0 and group_n > 0: + a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm + offs_bsn = offs_bn // group_n + b_scale_ptrs = ( + b_scale_ptr + off_experts * stride_bse + offs_bsn * stride_bsn + ) + # channel-wise + elif per_channel_quant: + b_scale_ptrs = ( + b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + ) + b_scale = tl.load(b_scale_ptrs) + # Load per-token scale for activations + a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm + a_scale = tl.load(a_scale_ptrs, mask=token_mask, other=0.0)[:, None] + # tensor-wise + else: + a_scale = tl.load(a_scale_ptr) + b_scale = tl.load(b_scale_ptr + off_experts) + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the + # K dimension. + if even_Ks: + a = tl.load( + a_ptrs, + mask=token_mask[:, None], + other=0.0, + ) + b = tl.load(b_ptrs) + else: + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (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) + + # We accumulate along the K dimension. + if use_int8_w8a16: + accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) + elif use_fp8_w8a8 or use_int8_w8a8: + if group_k > 0 and group_n > 0: + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_scale = tl.load( + a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0 + ) + b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) + + accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] + else: + if use_fp8_w8a8: + accumulator = tl.dot(a, b, acc=accumulator) + else: + accumulator += tl.dot(a, b) + else: + accumulator += tl.dot(a, b) + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if use_int8_w8a16: + accumulator *= b_scale + elif use_fp8_w8a8 or use_int8_w8a8: + if group_k == 0 or group_n == 0: + accumulator *= a_scale * b_scale + + if bias_ptr is not None: + accumulator += bias + + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0) + accumulator *= moe_weight[:, None] + + accumulator = accumulator.to(compute_type) + # ----------------------------------------------------------- + # Write back the block of the output + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +def invoke_fused_moe_kernel( + A: torch.Tensor, + B: torch.Tensor, + bias: Optional[torch.Tensor], + C: torch.Tensor, + A_scale: Optional[torch.Tensor], + B_scale: Optional[torch.Tensor], + B_zp: Optional[torch.Tensor], + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + mul_routed_weight: bool, + top_k: int, + config: Dict[str, Any], + compute_type: tl.dtype, + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + use_int4_w4a16: bool, + per_channel_quant: bool, + block_shape: Optional[List[int]] = None, + no_combine: bool = False, +) -> None: + assert topk_weights.stride(1) == 1 + assert sorted_token_ids.stride(0) == 1 + + padded_size = 0 + if use_fp8_w8a8: + assert B_scale is not None + if block_shape is None: + # activation tensor-wise fp8 quantization, dynamic or static + padded_size = padding_size + # activations apply per-token quantization when weights apply per-channel quantization by default + A, A_scale = scaled_fp8_quant( + A, A_scale, use_per_token_if_dynamic=per_channel_quant + ) + else: + # activation block-wise fp8 quantization + assert len(block_shape) == 2 + block_n, block_k = block_shape[0], block_shape[1] + if _is_cuda: + A, A_scale = sglang_per_token_group_quant_fp8(A, block_k) + else: + A, A_scale = per_token_group_quant_fp8(A, block_k) + assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1] + assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2] + assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1] + elif use_int8_w8a8: + assert B_scale is not None + if block_shape is None: + # activation channel-wise int8 quantization + assert ( + per_channel_quant + ), "int8 quantization only supports channel-wise quantization except for block-wise quantization" + A, A_scale = per_token_quant_int8(A) + else: + # activation block-wise int8 quantization + assert len(block_shape) == 2 + block_n, block_k = block_shape[0], block_shape[1] + if _is_cuda: + A, A_scale = sglang_per_token_group_quant_int8(A, block_k) + else: + A, A_scale = per_token_group_quant_int8(A, block_k) + assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1] + assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2] + assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1] + elif use_int8_w8a16 or use_int4_w4a16: + assert B_scale is not None + assert block_shape is None or block_shape[0] == 0 + else: + assert A_scale is None + assert B_scale is None + + grid = lambda META: ( + triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]) + * triton.cdiv(B.shape[1], META["BLOCK_SIZE_N"]), + ) + + K = B.shape[2] - padded_size + if K % config["BLOCK_SIZE_K"] == 0: + even_Ks = True + else: + even_Ks = False + + if ( + (use_int8_w8a16 or use_int4_w4a16) + and block_shape is not None + and block_shape[1] > 0 + ): + assert B_scale is not None and B_scale.ndim == 3 + assert B_zp is None or B_zp.ndim == 3 + assert bias is None + fused_moe_kernel_gptq_awq[grid]( + A, + B, + C, + B_scale, + B_zp, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + B.shape[1], + A.shape[1], + sorted_token_ids.shape[0], + topk_ids.numel(), + A.stride(0), + A.stride(1), + B.stride(0), + B.stride(2), + B.stride(1), + C.stride(1), + C.stride(2), + B_scale.stride(0), + B_scale.stride(2), + B_scale.stride(1), + B_zp.stride(0) if B_zp is not None else 0, + B_zp.stride(2) if B_zp is not None else 0, + B_zp.stride(1) if B_zp is not None else 0, + group_size=block_shape[1], + MUL_ROUTED_WEIGHT=mul_routed_weight, + top_k=top_k, + compute_type=compute_type, + has_zp=B_zp is not None, + use_int4_w4a16=use_int4_w4a16, + use_int8_w8a16=use_int8_w8a16, + even_Ks=even_Ks, + **config, + ) + + else: + + fused_moe_kernel[grid]( + A, + B, + bias, + C, + A_scale, + B_scale, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + B.shape[1], + B.shape[2] - padded_size, + sorted_token_ids.shape[0], + topk_ids.numel(), + A.stride(0), + A.stride(1), + B.stride(0), + B.stride(2), + B.stride(1), + bias.stride(0) if bias is not None else 0, + bias.stride(1) if bias is not None else 0, + C.stride(1), + C.stride(2), + A_scale.stride(0) if A_scale is not None and A_scale.ndim == 2 else 0, + A_scale.stride(1) if A_scale is not None and A_scale.ndim == 2 else 0, + B_scale.stride(0) if B_scale is not None and B_scale.ndim >= 2 else 0, + B_scale.stride(2) if B_scale is not None and B_scale.ndim == 3 else 0, + B_scale.stride(1) if B_scale is not None and B_scale.ndim >= 2 else 0, + 0 if block_shape is None else block_shape[0], + 0 if block_shape is None else block_shape[1], + MUL_ROUTED_WEIGHT=mul_routed_weight, + top_k=top_k, + compute_type=compute_type, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + per_channel_quant=per_channel_quant, + even_Ks=even_Ks, + **config, + ) + + +# _moe_sum_reduce_kernel kernel modified from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/moe_sum_reduce.py +@triton.jit +def _moe_sum_reduce_kernel( + input_ptr, + input_stride_0, + input_stride_1, + input_stride_2, + output_ptr, + output_stride_0, + output_stride_1, + token_num: int, + topk_num: int, + hidden_dim: int, + routed_scaling_factor: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DIM: tl.constexpr, + NUM_STAGE: tl.constexpr, +): + input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64) + input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64) + output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64) + + token_block_id = tl.program_id(0) + dim_block_id = tl.program_id(1) + + offs_token = token_block_id * BLOCK_M + tl.arange(0, BLOCK_M) + offs_dim = dim_block_id * BLOCK_DIM + tl.arange(0, BLOCK_DIM) + + mask_token = offs_token < token_num + mask_dim = offs_dim < hidden_dim + + base_ptrs = input_ptr + offs_token[:, None] * input_stride_0 + offs_dim[None, :] + + accumulator = tl.zeros((BLOCK_M, BLOCK_DIM), dtype=tl.float32) + + for i in tl.range(0, topk_num, num_stages=NUM_STAGE): + tile = tl.load( + base_ptrs + i * input_stride_1, + mask=mask_token[:, None] & mask_dim[None, :], + other=0.0, + ) + accumulator += tile.to(tl.float32) + accumulator *= routed_scaling_factor + + # -------- Write back -------- + store_ptrs = output_ptr + offs_token[:, None] * output_stride_0 + offs_dim[None, :] + tl.store( + store_ptrs, + accumulator.to(input_ptr.dtype.element_ty), + mask=mask_token[:, None] & mask_dim[None, :], + ) + + +def moe_sum_reduce_triton( + input: torch.Tensor, output: torch.Tensor, routed_scaling_factor: float +): + assert input.is_contiguous() + assert output.is_contiguous() + + token_num, topk_num, hidden_dim = input.shape + assert output.shape[0] == token_num and output.shape[1] == hidden_dim + + BLOCK_M = 1 + BLOCK_DIM = 2048 + NUM_STAGE = 1 + num_warps = 16 + + grid = ( + triton.cdiv(token_num, BLOCK_M), + triton.cdiv(hidden_dim, BLOCK_DIM), + ) + + _moe_sum_reduce_kernel[grid]( + input, + *input.stride(), + output, + *output.stride(), + token_num=token_num, + topk_num=topk_num, + hidden_dim=hidden_dim, + routed_scaling_factor=routed_scaling_factor, + BLOCK_M=BLOCK_M, + BLOCK_DIM=BLOCK_DIM, + NUM_STAGE=NUM_STAGE, + num_warps=num_warps, + ) + return diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py new file mode 100644 index 000000000..0ea1fa1eb --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -0,0 +1,1078 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/a6221a144af772fd1a68fe7e627935dc53e81738/vllm/model_executor/layers/fused_moe/layer.py + +import logging +from enum import Enum +from typing import List, Optional, Tuple + +import torch + +from sglang.srt.distributed import ( + get_moe_expert_parallel_rank, + get_moe_expert_parallel_world_size, + get_moe_tensor_parallel_rank, + get_moe_tensor_parallel_world_size, + get_tp_group, + tensor_model_parallel_all_reduce, +) +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.eplb.expert_location import get_global_expert_location_metadata +from sglang.srt.layers.moe import ( + MoeRunnerConfig, + get_moe_runner_backend, + should_use_flashinfer_trtllm_moe, +) +from sglang.srt.layers.moe.token_dispatcher.standard import ( + CombineInput, + StandardDispatcher, +) +from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod +from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4FusedMoEMethod +from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_loader.weight_utils import narrow_padded_param_and_loaded_weight +from sglang.srt.utils import ( + cpu_has_amx_support, + get_bool_env_var, + is_cpu, + is_flashinfer_available, + is_hip, + next_power_of_2, + round_up, +) + +if is_flashinfer_available(): + from flashinfer import ( + RoutingMethodType, + fp4_quantize, + reorder_rows_for_gated_act_gemm, + shuffle_matrix_a, + shuffle_matrix_sf_a, + ) + +_is_hip = is_hip() +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() + + +# Try to import FP4 TRTLLM function if flashinfer is available +trtllm_fp4_block_scale_moe = None +if should_use_flashinfer_trtllm_moe(): + try: + from flashinfer.fused_moe import trtllm_fp4_block_scale_moe + except ImportError: + trtllm_fp4_block_scale_moe = None + +logger = logging.getLogger(__name__) + + +def _is_fp4_quantization_enabled(): + """Check if ModelOpt FP4 quantization is enabled.""" + try: + # Use the same simple check that works for class selection + quantization = global_server_args_dict.get("quantization") + return quantization == "modelopt_fp4" + except: + return False + + +def _get_tile_tokens_dim(num_tokens, top_k, num_experts): + # Guess tokens per expert assuming perfect expert distribution first. + num_tokens_per_expert = (num_tokens * top_k) // num_experts + # And pad the number to the next power of 2. + tile_tokens_dim = next_power_of_2(num_tokens_per_expert) + # Cap to 8-64 tokens per CTA tile as it's the range supported by the kernel. + tile_tokens_dim = min(max(tile_tokens_dim, 8), 64) + return tile_tokens_dim + + +class FusedMoeWeightScaleSupported(Enum): + TENSOR = "tensor" + CHANNEL = "channel" + GROUP = "group" + BLOCK = "block" + + +class FusedMoE(torch.nn.Module): + """FusedMoE layer for MoE models. + + This layer contains both MergedColumnParallel weights (gate_up_proj / + w13) and RowParallelLinear weights (down_proj/ w2). + + Note: Mixtral uses w1, w2, and w3 for gate, up, and down_proj. We + copy that naming convention here and handle any remapping in the + load_weights function in each model implementation. + + Args: + num_experts: Number of experts in the model + top_k: Number of experts selected for each token + hidden_size: Input hidden state size of the transformer + intermediate_size: Intermediate size of the experts + params_dtype: Data type for the parameters. + reduce_results: Whether to apply all_reduce on the output of the layer + quant_config: Quantization configuration. + inplace: suggestion to compute inplace (modify input activation). + """ + + def __init__( + self, + num_experts: int, + hidden_size: int, + intermediate_size: int, + layer_id: int, + top_k: Optional[int] = None, + num_fused_shared_experts: int = 0, + params_dtype: Optional[torch.dtype] = None, + reduce_results: bool = False, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + activation: str = "silu", + apply_router_weight_on_input: bool = False, + use_presharded_weights: bool = False, + inplace: bool = True, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + gemm1_alpha: Optional[float] = None, + gemm1_clamp_limit: Optional[float] = None, + use_weight_loader_fused: bool = False, + with_bias=False, + ): + super().__init__() + + if params_dtype is None: + params_dtype = torch.get_default_dtype() + + self.layer_id = layer_id + self.top_k = top_k + self.hidden_size = hidden_size + self.num_experts = num_experts + self.num_fused_shared_experts = num_fused_shared_experts + self.expert_map_cpu = None + self.expert_map_gpu = None + + enable_flashinfer_cutlass_moe = get_moe_runner_backend().is_flashinfer_cutlass() + + if enable_flashinfer_cutlass_moe and quant_config is None: + logger.warning("Disable flashinfer MoE when quantization config is None.") + enable_flashinfer_cutlass_moe = False + + self.enable_flashinfer_cutlass_moe = enable_flashinfer_cutlass_moe + self.moe_ep_size = get_moe_expert_parallel_world_size() + self.moe_ep_rank = get_moe_expert_parallel_rank() + self.moe_tp_size = get_moe_tensor_parallel_world_size() + self.moe_tp_rank = get_moe_tensor_parallel_rank() + assert num_experts % self.moe_ep_size == 0 + self.num_local_experts = num_experts // self.moe_ep_size + self.start_expert_id = self.moe_ep_rank * self.num_local_experts + self.end_expert_id = self.start_expert_id + self.num_local_experts - 1 + if self.moe_ep_size > 1: + # TODO(ch-wan): support shared experts fusion + # Create a tensor of size num_experts filled with -1 + self.expert_map_cpu = torch.full( + (self.num_experts,), -1, dtype=torch.int32, device="cpu" + ) + # Create a expert map for the local experts + self.expert_map_cpu[ + self.moe_ep_rank + * self.num_local_experts : (self.moe_ep_rank + 1) + * self.num_local_experts + ] = torch.arange(0, self.num_local_experts, dtype=torch.int32, device="cpu") + + assert intermediate_size % self.moe_tp_size == 0 + self.intermediate_size_per_partition = intermediate_size // self.moe_tp_size + self.reduce_results = reduce_results + self.use_presharded_weights = use_presharded_weights + + self.use_triton_kernels = get_moe_runner_backend().is_triton_kernel() + + self.quant_config = quant_config + self.use_flashinfer_mxfp4_moe = get_moe_runner_backend().is_flashinfer_mxfp4() + # TODO maybe we should remove this `if`, since `Mxfp4MoEMethod` does another round-up logic + if ( + self.quant_config is not None + and self.quant_config.get_name() == "mxfp4" + and self.use_flashinfer_mxfp4_moe + ): + hidden_size = round_up(hidden_size, 256) + self.hidden_size = hidden_size + + self.moe_runner_config = MoeRunnerConfig( + num_experts=num_experts, + num_local_experts=self.num_local_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=self.intermediate_size_per_partition, + layer_id=layer_id, + top_k=top_k, + num_fused_shared_experts=num_fused_shared_experts, + params_dtype=params_dtype, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + inplace=inplace, + no_combine=no_combine, + routed_scaling_factor=routed_scaling_factor, + gemm1_alpha=gemm1_alpha, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + if quant_config is None: + self.quant_method: FusedMoEMethodBase = UnquantizedFusedMoEMethod( + self.use_triton_kernels + ) + else: + self.quant_method: FusedMoEMethodBase = quant_config.get_quant_method( + self, prefix + ) + assert self.quant_method is not None + + self.quant_method.create_weights( + layer=self, + num_experts=self.num_local_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=self.intermediate_size_per_partition, + params_dtype=params_dtype, + weight_loader=( + self.weight_loader + if not use_weight_loader_fused + else self.weight_loader_fused + ), + with_bias=with_bias, + ) + + self.quant_method.create_moe_runner(self, self.moe_runner_config) + self.dispatcher = StandardDispatcher() + + def _load_per_tensor_weight_scale( + self, + shard_id: str, + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + expert_id: int, + ): + param_data = param.data + # for per tensor weight quantization + if shard_id in ("w1", "w3"): + # We have to keep the weight scales of w1 and w3 because + # we need to re-quantize w1/w3 weights after weight loading. + idx = 0 if shard_id == "w1" else 1 + param_data[expert_id][idx] = loaded_weight + # If we are in the row parallel case (down_proj) + elif shard_id == "w2": + param_data[expert_id] = loaded_weight + + def _load_model_weight_or_group_weight_scale( + self, + shard_dim: int, + expert_data: torch.Tensor, + shard_id: str, + loaded_weight: torch.Tensor, + tp_rank: int, + is_bias: bool = False, + ): + # Load grouped weight scales for group quantization + # or model weights + if shard_id == "w2": + self._load_w2( + shard_id=shard_id, + shard_dim=shard_dim, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + is_bias=is_bias, + ) + elif shard_id in ("w1", "w3", "w13"): + self._load_w13( + shard_id=shard_id, + shard_dim=shard_dim, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + is_bias=is_bias, + ) + + def _load_per_channel_weight_scale( + self, + expert_data: torch.Tensor, + shard_dim: int, + shard_id: str, + loaded_weight: torch.Tensor, + tp_rank: int, + ): + # for per channel weight quantization + if shard_id == "w2": + expert_data.copy_(loaded_weight) + elif shard_id in ("w1", "w3"): + self._load_w13( + shard_id=shard_id, + shard_dim=shard_dim, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + ) + + def _load_w13( + self, + expert_data: torch.Tensor, + shard_dim: int, + shard_id: str, + loaded_weight: torch.Tensor, + tp_rank: int, + is_bias: bool = False, + ): + + # Index the loaded weight for tp sharding. + # gate_up_proj: "MergedColumnParallel", so tp sharding on output_dim + assert shard_id in {"w1", "w3", "w13"} + + if is_bias: + # if this weight is a bias, the last dimension must be the sharded dimension + shard_dim = -1 + + if shard_id in {"w1", "w3"}: + # non-fused version + shard_size = expert_data.shape[shard_dim] // 2 + elif shard_id in {"w13"}: + # fused version + shard_size = expert_data.shape[shard_dim] + else: + raise NotImplementedError + + # Narrow parameter and load. + # w1, gate_proj: Load into first logical weight of w13. + # w3, up_proj: Load into second logical weight of w13. + # trtllm cutlass kernel assumes differently + switch_w13 = getattr(self.quant_method, "load_up_proj_weight_first", False) + if (switch_w13 and shard_id == "w1") or (not switch_w13 and shard_id == "w3"): + start = shard_size + else: + start = 0 + + if _is_cpu: + expert_data, loaded_weight = narrow_padded_param_and_loaded_weight( + expert_data, + loaded_weight, + start, + shard_size * tp_rank, + shard_dim, + shard_size, + not self.use_presharded_weights, + ) + else: + if not self.use_presharded_weights: + if not is_bias and self.use_triton_kernels: + # do not transpose for bias + loaded_weight = loaded_weight.transpose(-2, -1) + loaded_weight = loaded_weight.narrow( + shard_dim, shard_size * tp_rank, shard_size + ) + + expert_data = expert_data.narrow(shard_dim, start, shard_size) + expert_data.copy_(loaded_weight) + + def _load_w2( + self, + expert_data: torch.Tensor, + shard_dim: int, + shard_id: str, + loaded_weight: torch.Tensor, + tp_rank: int, + is_bias: bool = False, + ): + """Load w2 weights for down projection. + + Args: + expert_data: The expert data tensor to load into + shard_dim: The dimension to shard along + shard_id: The shard ID (must be "w2") + loaded_weight: The weight tensor to load from + tp_rank: The tensor parallel rank + """ + if not isinstance(expert_data, torch.Tensor) or not isinstance( + loaded_weight, torch.Tensor + ): + raise ValueError("expert_data and loaded_weight must be torch.Tensor") + + if ( + self.quant_config is not None + and "modelopt" in self.quant_config.get_name() + and (expert_data.dim() != 2 or loaded_weight.dim() != 2) + ): + raise ValueError( + f"Expected 2D tensors, got expert_data shape {expert_data.shape} and loaded_weight shape {loaded_weight.shape}" + ) + + if shard_id != "w2": + raise ValueError(f"shard_id must be 'w2', got {shard_id}") + + # Index the loaded weight for tp sharding. + # down_proj: "RowParallel" so tp sharding on input_dim + # Narrow parameter and load. + if is_bias: + # this expert_data is a bias, not weight, + # for w2_weight_bias in TP, it does not need to be sharded + shard_size = expert_data.shape[-1] + else: + # this parameter is a weight matrix + # for w2 in TP, it shards the input_features, i.e., shard_dim=2 + shard_size = expert_data.shape[shard_dim] + + if _is_cpu: + expert_data, loaded_weight = narrow_padded_param_and_loaded_weight( + expert_data, + loaded_weight, + 0, # param_data_start + shard_size * tp_rank, + shard_dim, + shard_size, + not self.use_presharded_weights, + ) + else: + if not is_bias and not self.use_presharded_weights: + if self.use_triton_kernels: + loaded_weight = loaded_weight.transpose(-2, -1) + loaded_weight = loaded_weight.narrow( + shard_dim, shard_size * tp_rank, shard_size + ) + + # w2, down_proj: Load into only logical weight of w2. + expert_data.copy_(loaded_weight) + + def _load_single_value( + self, param: torch.nn.Parameter, loaded_weight: torch.Tensor, expert_id: int + ): + param_data = param.data + + # Input scales can be loaded directly and should be equal. + param_data[expert_id] = loaded_weight + + def _load_g_idx( + self, + shard_id: str, + expert_data: torch.Tensor, + shard_dim: int, + loaded_weight: torch.Tensor, + tp_rank: int, + ): + + if shard_id == "w2": + self._load_w2( + shard_id=shard_id, + shard_dim=shard_dim, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + ) + else: + assert shard_id in ("w1", "w3") + expert_data.copy_(loaded_weight) + + def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int: + if self.expert_map_cpu is None: + return expert_id + return self.expert_map_cpu[expert_id].item() + + def weight_loader( + self, + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: Optional[int], + ) -> None: + + # if expert_id is None, then + # all the experts are loaded at the same time + if ( + not expert_id + and self.quant_config is not None + and self.quant_config.get_name() == "mxfp4" + and self.quant_config.is_static_cfg() + ): + if "bias" in weight_name: + dim1 = loaded_weight.shape[1] + param.data[:, :dim1].copy_(loaded_weight) + else: + dim1 = loaded_weight.shape[1] + dim2 = loaded_weight.shape[2] + param.data[:, :dim1, :dim2].copy_(loaded_weight) + return + + global_expert_location_metadata = get_global_expert_location_metadata() + if global_expert_location_metadata is None: + self._weight_loader_impl( + param=param, + loaded_weight=loaded_weight, + weight_name=weight_name, + shard_id=shard_id, + expert_id=expert_id, + ) + return + + if expert_id >= self.num_experts - self.num_fused_shared_experts: + # This is a shared expert. + physical_expert_ids = [expert_id] + else: + physical_expert_ids = ( + global_expert_location_metadata.logical_to_all_physical( + self.layer_id, expert_id + ) + ) + + for physical_expert_id in physical_expert_ids: + self._weight_loader_physical( + param=param, + loaded_weight=loaded_weight, + weight_name=weight_name, + shard_id=shard_id, + expert_id=physical_expert_id, + ) + + def _weight_loader_physical( + self, + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + ) -> None: + + expert_id = self._map_global_expert_id_to_local_expert_id(expert_id) + if expert_id == -1: + return + self._weight_loader_impl( + param=param, + loaded_weight=loaded_weight, + weight_name=weight_name, + shard_id=shard_id, + expert_id=expert_id, + ) + + def _weight_loader_impl( + self, + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + ) -> None: + + tp_rank = self.moe_tp_rank + + # compressed-tensors checkpoints with packed weights are stored flipped + # TODO (mgoin): check self.quant_method.quant_config.quant_format + # against known CompressionFormat enum values that have this quality + loaded_weight = ( + loaded_weight.t().contiguous() + if ( + self.quant_method.__class__.__name__ + == "CompressedTensorsWNA16MoEMethod" + ) + else loaded_weight + ) + + if shard_id not in ("w1", "w2", "w3"): + raise ValueError( + f"shard_id must be ['w1','w2','w3'] but " f"got {shard_id}." + ) + + # Flashinfer assumes w31 format for w13_weight. Same for the scales. + if should_use_flashinfer_trtllm_moe(): + shard_id = {"w1": "w3", "w3": "w1", "w2": "w2"}[shard_id] + + WEIGHT_SCALE_SUPPORTED = [e.value for e in FusedMoeWeightScaleSupported] + # Fetch the dim to shard the parameter/loaded weight + # based on the shard id. This will be whatever + # dimension intermediate_size is used. + SHARD_ID_TO_SHARDED_DIM = {"w1": 0, "w2": 1, "w3": 0} + + expert_data = param.data[expert_id] + + # is_transposed: if the dim to shard the weight + # should be flipped. Required by GPTQ, compressed-tensors + # should be whatever dimension intermediate_size is + is_transposed = getattr(param, "is_transposed", False) + shard_dim = SHARD_ID_TO_SHARDED_DIM[shard_id] + if self.use_triton_kernels: + is_transposed = True + if is_transposed: + shard_dim = int(not shard_dim) + + # Case input scale: input_scale loading is only supported for fp8 + if "input_scale" in weight_name: + # INT4-FP8 (INT4 MoE Weight, FP8 Compute): Adjust input_scale for e4m3fnuz (AMD) + if _is_hip and get_bool_env_var("SGLANG_INT4_WEIGHT"): + loaded_weight = loaded_weight * 2.0 + + # this is needed for compressed-tensors only + loaded_weight = loaded_weight.to(param.data.device) + + if ( + ( + "compressed" in self.quant_method.__class__.__name__.lower() + or "w4afp8" in self.quant_config.get_name() + ) + and (param.data[expert_id] != 1).any() + and ((param.data[expert_id] - loaded_weight).abs() > 1e-5).any() + ): + raise ValueError( + "input_scales of w1 and w3 of a layer " + f"must be equal. But got {param.data[expert_id]} " + f"vs. {loaded_weight}" + ) + + self._load_single_value( + param=param, loaded_weight=loaded_weight, expert_id=expert_id + ) + return + + # Case g_idx + if "g_idx" in weight_name: + self._load_g_idx( + shard_dim=0, + shard_id=shard_id, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + ) + return + + if "ModelOpt" in self.quant_method.__class__.__name__: + # Determine per-tensor weight scale patterns based on variant + is_fp4_variant = isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod) + + # FP4 uses "weight_scale_2" for per-tensor, FP8 uses "weight_scale" for per-tensor + per_tensor_conditions = ( + "weight_scale_2" in weight_name + if is_fp4_variant + else "weight_scale" in weight_name + ) or "input_scale" in weight_name + + if per_tensor_conditions: + self._load_per_tensor_weight_scale( + shard_id=shard_id, + param=param, + loaded_weight=loaded_weight, + expert_id=expert_id, + ) + elif "weight" in weight_name: + self._load_model_weight_or_group_weight_scale( + shard_id=shard_id, + shard_dim=shard_dim, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + ) + return + + # Case weight scales and zero_points + if "scale" in weight_name or "zero" in weight_name or "offset" in weight_name: + # load the weight scales and zp based on the quantization scheme + # supported weight scales/zp can be found in + # FusedMoeWeightScaleSupported + # TODO @dsikka: once hardened, refactor to use vLLM Parameters + # specific to each case + quant_method = getattr(param, "quant_method", None) + if quant_method == FusedMoeWeightScaleSupported.CHANNEL.value: + # INT4-FP8 (INT4 MoE Weight, FP8 Compute): Adjust INT4 column-wise scaling number to e4m3fnuz (AMD) + if _is_hip and get_bool_env_var("SGLANG_INT4_WEIGHT"): + loaded_weight = loaded_weight * 0.5 + + self._load_per_channel_weight_scale( + shard_id=shard_id, + shard_dim=shard_dim, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + ) + elif quant_method in [ + FusedMoeWeightScaleSupported.GROUP.value, + FusedMoeWeightScaleSupported.BLOCK.value, + ]: + self._load_model_weight_or_group_weight_scale( + shard_id=shard_id, + shard_dim=shard_dim, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + ) + elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value: + # INT4-FP8 (INT4 MoE Weight, FP8 Compute): Adjust FP8 per-tensor scaling number for e4m3fnuz (AMD) + if _is_hip and get_bool_env_var("SGLANG_INT4_WEIGHT"): + loaded_weight = loaded_weight * 2.0 + + self._load_per_tensor_weight_scale( + shard_id=shard_id, + param=param, + loaded_weight=loaded_weight, + expert_id=expert_id, + ) + else: + raise ValueError( + f"quant method must be one of {WEIGHT_SCALE_SUPPORTED}" + ) + return + + # Case weight_shape + if "weight_shape" in weight_name: + # only required by compressed-tensors + self._load_single_value( + param=param, loaded_weight=loaded_weight, expert_id=expert_id + ) + return + + # Case model weights + if "weight" in weight_name: + self._load_model_weight_or_group_weight_scale( + shard_id=shard_id, + shard_dim=shard_dim, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + ) + return + + def weight_loader_fused( + self, + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + ) -> None: + tp_rank = self.moe_tp_rank + + if ( + self.quant_config is not None + and self.quant_config.get_name() == "mxfp4" + and self.quant_config.is_static_cfg() + ): + if "bias" in weight_name: + dim1 = loaded_weight.shape[1] + param.data[:, :dim1].copy_(loaded_weight) + elif "scale" in weight_name: + param.data.copy_(loaded_weight) + else: + dim1 = loaded_weight.shape[1] + dim2 = loaded_weight.shape[2] + param.data[:, :dim1, :dim2].copy_(loaded_weight) + return + + # compressed-tensors checkpoints with packed weights are stored flipped + # TODO: check self.quant_method.quant_config.quant_format + # against known CompressionFormat enum values that have this quality + loaded_weight = ( + loaded_weight.t().contiguous() + if ( + self.quant_method.__class__.__name__ + == "CompressedTensorsWNA16MoEMethod" + ) + else loaded_weight + ) + + if shard_id not in ("w13", "w2"): + raise ValueError(f"shard_id must be ['w13','w2'] but " f"got {shard_id}.") + + # Fetch the dim to shard the parameter/loaded weight + # based on the shard id. This will be whatever + # dimension intermediate_size is used. + SHARD_ID_TO_SHARDED_DIM = {"w13": 1, "w2": 2} + SHARD_ID_TO_SHARDED_DIM_TRANSPOSE = {"w13": 2, "w2": 1} + + expert_data = param.data + is_bias = expert_data.dim() == 2 + + # is_transposed: if the dim to shard the weight + # should be flipped. Required by GPTQ, compressed-tensors + # should be whatever dimension intermediate_size is + is_transposed = getattr(param, "is_transposed", False) + + if self.use_triton_kernels: + is_transposed = True + shard_dim = ( + SHARD_ID_TO_SHARDED_DIM[shard_id] + if not is_transposed + else SHARD_ID_TO_SHARDED_DIM_TRANSPOSE[shard_id] + ) + + # Case model weights + if "weight" in weight_name: + self._load_model_weight_or_group_weight_scale( + shard_id=shard_id, + shard_dim=shard_dim, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=tp_rank, + is_bias=is_bias, + ) + return + else: + logging.warning( + f"Unsupported weight_name {weight_name} for FusedMoE weight_loader_fused. Nothing is loaded." + ) + + def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput): + origin_hidden_states_dim = hidden_states.shape[-1] + assert self.quant_method is not None + + if self.moe_ep_size > 1 and not self.enable_flashinfer_cutlass_moe: + if self.expert_map_cpu is not None and self.expert_map_gpu is None: + # If we are in EP mode, we need to move the expert map to GPU. + self.expert_map_gpu = self.expert_map_cpu.to(device="cuda") + + if self.expert_map_gpu is not None: + if TopKOutputChecker.format_is_standard(topk_output): + topk_output = topk_output._replace( + topk_ids=self.expert_map_gpu[topk_output.topk_ids] + ) + elif TopKOutputChecker.format_is_triton_kernel(topk_output): + raise NotImplementedError() + + dispatch_output = self.dispatcher.dispatch( + hidden_states=hidden_states, topk_output=topk_output + ) + + # TODO: consider using symmetric memory + combine_input = self.quant_method.apply( + layer=self, + dispatch_output=dispatch_output, + ) + + final_hidden_states = self.dispatcher.combine(combine_input) + + final_hidden_states = final_hidden_states[ + ..., :origin_hidden_states_dim + ].contiguous() + + if self.reduce_results and (self.moe_tp_size > 1 or self.moe_ep_size > 1): + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states + + @classmethod + def make_expert_params_mapping( + cls, + ckpt_gate_proj_name: str, + ckpt_down_proj_name: str, + ckpt_up_proj_name: str, + num_experts: int, + ) -> List[Tuple[str, str, int, str]]: + + return [ + # (param_name, weight_name, expert_id, shard_id) + ( + ( + "experts.w13_" + if weight_name in [ckpt_gate_proj_name, ckpt_up_proj_name] + else "experts.w2_" + ), + f"experts.{expert_id}.{weight_name}.", + expert_id, + shard_id, + ) + for expert_id in range(num_experts) + for shard_id, weight_name in [ + ("w1", ckpt_gate_proj_name), + ("w2", ckpt_down_proj_name), + ("w3", ckpt_up_proj_name), + ] + ] + + @classmethod + def make_expert_params_mapping_fused( + cls, + ckpt_gate_up_proj_name: str, + ckpt_down_proj_name: str, + ckpt_gate_up_proj_bias_name: str, + ckpt_down_proj_bias_name: str, + ): + return [ + ("experts.w13_weight", f"experts.{ckpt_gate_up_proj_name}", "w13"), + ( + "experts.w13_weight_bias", + f"experts.{ckpt_gate_up_proj_bias_name}", + "w13", + ), + ("experts.w2_weight", f"experts.{ckpt_down_proj_name}", "w2"), + ("experts.w2_weight_bias", f"experts.{ckpt_down_proj_bias_name}", "w2"), + ] + + @classmethod + def make_expert_params_mapping_fused_mxfp4( + cls, + ckpt_gate_up_proj_name: str, + ckpt_down_proj_name: str, + ckpt_gate_up_proj_bias_name: str, + ckpt_down_proj_bias_name: str, + ckpt_gate_up_proj_scale_name: str, + ckpt_down_proj_scale_name: str, + ): + return [ + ("experts.w13_weight", f"experts.{ckpt_gate_up_proj_name}", "w13"), + ( + "experts.w13_weight_bias", + f"experts.{ckpt_gate_up_proj_bias_name}", + "w13", + ), + ("experts.w2_weight", f"experts.{ckpt_down_proj_name}", "w2"), + ("experts.w2_weight_bias", f"experts.{ckpt_down_proj_bias_name}", "w2"), + ( + "experts.w13_weight_scale", + f"experts.{ckpt_gate_up_proj_scale_name}", + "w13", + ), + ("experts.w2_weight_scale", f"experts.{ckpt_down_proj_scale_name}", "w2"), + ] + + @classmethod + def make_expert_input_scale_params_mapping( + cls, + num_experts: int, + ) -> List[Tuple[str, str, int, str]]: + # (param_name, weight_name, expert_id, shard_id) + return [ + ( + "experts.w13_" if shard_id in ["w1", "w3"] else "experts.w2_", + f"experts.{expert_id}.{shard_id}.", + expert_id, + shard_id, + ) + for expert_id in range(num_experts) + for shard_id in ["w1", "w2", "w3"] + ] + + def should_fuse_routed_scaling_factor_in_topk(self): + return isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod) or ( + isinstance(self.quant_method, Fp8MoEMethod) + and self.quant_method.use_cutlass_fused_experts_fp8 + ) + + +class FlashInferFusedMoE(FusedMoE): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.use_flashinfer_trtllm_moe = should_use_flashinfer_trtllm_moe() + + def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput): + assert self.use_flashinfer_trtllm_moe + assert ( + self.moe_runner_config.activation == "silu" + ), "Only silu is supported for flashinfer blockscale fp8 moe" + assert self.quant_method is not None + assert ( + topk_output.topk_config.renormalize + ), "Renormalize is required for flashinfer blockscale fp8 moe" + assert ( + self.num_fused_shared_experts == 0 + ), "Fused shared experts are not supported for flashinfer blockscale fp8 moe" + + assert TopKOutputChecker.format_is_bypassed(topk_output) + + # Matrix multiply. + final_hidden_states = self.quant_method.apply_with_router_logits( + layer=self, + x=hidden_states, + topk_output=topk_output, + ) + + if self.reduce_results and (self.moe_tp_size > 1 or self.moe_ep_size > 1): + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states + + +class FlashInferFP4MoE(FusedMoE): + """FP4 TRTLLM MoE implementation using FlashInfer.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # --------------------------------------------------------------------- + # Helper: quantize hidden states to FP4 each forward pass + # --------------------------------------------------------------------- + def _quantize_hidden_states_fp4(self, hidden_states: torch.Tensor): + """ + Quantize hidden states using global scale factor from quantization method. + + Global scale factor is set by ModelOptNvFp4FusedMoEMethod during weight loading. + Only block scales are computed at runtime for efficiency. + + Returns (packed_fp4_uint8, scale_float8_e4m3fn_runtime, global_scale_float32) + """ + + # flashinfer.fp4_quantize returns (packed_uint8, scale_fp8) + # Only the block scales are computed at runtime + hs_fp4_bytes, hs_sf_bytes = fp4_quantize( + hidden_states, + self.w13_input_scale_quant, + 16, # sf_vec_size + False, # use_ue8m0 + False, # is_sf_swizzled_layout + ) + + hs_fp4 = hs_fp4_bytes.reshape( + hidden_states.shape[0], hidden_states.shape[1] // 2 + ) + hs_sf = hs_sf_bytes.view(torch.float8_e4m3fn).reshape(-1) + + return hs_fp4, hs_sf + + def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput): + """Forward pass using FP4 TRTLLM kernel. + + Args: + hidden_states: Input tensor + topk_output: TopKOutput object with Bypassed format + """ + assert isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod) + + assert TopKOutputChecker.format_is_bypassed(topk_output) + + router_logits = topk_output.router_logits + topk_config = topk_output.topk_config + + hs_fp4, hs_scale_linear = self._quantize_hidden_states_fp4(hidden_states) + + router_logits = router_logits.to(torch.float32) + + result = trtllm_fp4_block_scale_moe( + routing_logits=router_logits, + routing_bias=topk_config.correction_bias.to(hidden_states.dtype), + hidden_states=hs_fp4, + hidden_states_scale=hs_scale_linear.view(torch.float8_e4m3fn).flatten(), + gemm1_weights=self.gemm1_weights_fp4_shuffled.data, + gemm1_weights_scale=self.gemm1_scales_fp4_shuffled.data.view( + torch.float8_e4m3fn + ), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=self.gemm2_weights_fp4_shuffled.data, + gemm2_weights_scale=self.gemm2_scales_fp4_shuffled.data.view( + torch.float8_e4m3fn + ), + gemm2_bias=None, + output1_scale_scalar=self.g1_scale_c.data, + output1_scale_gate_scalar=self.g1_alphas.data, + output2_scale_scalar=self.g2_alphas.data, + num_experts=self.num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.moe_ep_rank * self.num_local_experts, + local_num_experts=self.num_local_experts, + routed_scaling_factor=self.moe_runner_config.routed_scaling_factor, + tile_tokens_dim=_get_tile_tokens_dim( + hidden_states.shape[0], topk_config.top_k, self.num_local_experts + ), + routing_method_type=RoutingMethodType.DeepSeekV3, + do_finalize=True, + )[0] + + return result diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py b/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py new file mode 100644 index 000000000..64d0126d6 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +from typing import Tuple + +import torch +import triton + +from sglang.srt.utils import is_cuda, is_hip + +_is_cuda = is_cuda() +_is_hip = is_hip() + +if _is_cuda or _is_hip: + from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size + + +def moe_align_block_size( + topk_ids: torch.Tensor, block_size: int, num_experts: int +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Aligns the token distribution across experts to be compatible with block + size for matrix multiplication. + + Parameters: + - topk_ids: A tensor of shape [total_tokens, top_k] representing the + top-k expert indices for each token. + - block_size: The block size used in block matrix multiplication. + - num_experts: The total number of experts. + + Returns: + - sorted_token_ids: A tensor containing the sorted token indices according + to their allocated expert. + - expert_ids: A tensor indicating the assigned expert index for each block. + - num_tokens_post_padded: The total number of tokens after padding, + ensuring divisibility by block_size. + + This function pads the number of tokens that each expert needs to process + so that it is divisible by block_size. + Padding ensures that during block matrix multiplication, the dimensions + align correctly. + + Example: + Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]], + block_size = 4, and num_experts = 4: + - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts, + with each expert needing to process 3 tokens. + - As block_size is 4, we pad 1 token for each expert. + - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3]. + - Then append padding tokens [12, 12, 12, 12] for each block. + - After sorting by expert index, we obtain token_ids + [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12]. + Tokens 12 are non-existent (padding) and are ignored in + the subsequent matrix multiplication. + - The padding ensures that the total number of tokens is now divisible + by block_size for proper block matrix operations. + """ + max_num_tokens_padded = topk_ids.numel() + (num_experts + 1) * (block_size - 1) + sorted_ids = torch.empty( + (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device + ) + max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size) + expert_ids = torch.empty( + (max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device + ) + num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device) + + # In EP, expert_ids for filtered experts are -1. We have num_experts + 1 ids in total. + cumsum_buffer = torch.empty( + (num_experts + 2,), dtype=torch.int32, device=topk_ids.device + ) + + # Threshold based on benchmark results + fuse_sorted_ids_padding = sorted_ids.shape[0] <= 4096 + if not fuse_sorted_ids_padding: + sorted_ids.fill_(topk_ids.numel()) + + sgl_moe_align_block_size( + topk_ids, + num_experts + 1, + block_size, + sorted_ids, + expert_ids, + num_tokens_post_pad, + cumsum_buffer, + fuse_sorted_ids_padding, + ) + return sorted_ids, expert_ids, num_tokens_post_pad diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py new file mode 100644 index 000000000..5d39b8bbc --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py @@ -0,0 +1,335 @@ +# Adapted from https://github.com/vllm-project/vllm/pull/18595/files#diff-f426a6de78c82ffec568eff6811bfbf0043dab5f87f1a8c0cffdbdcb8a81e035 + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch +from sgl_kernel import gelu_and_mul, silu_and_mul +from triton_kernels.matmul_ogs import ( + FlexCtx, + FnSpecs, + FusedActivation, + PrecisionConfig, + matmul_ogs, +) +from triton_kernels.numerics import InFlexData +from triton_kernels.routing import GatherIndx, RoutingData, ScatterIndx +from triton_kernels.swiglu import swiglu_fn + +if TYPE_CHECKING: + from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig + from sglang.srt.layers.moe.topk import TopKOutput + + +def quantize(w, dtype, dev, **opt): + if dtype == "bf16": + return w.to(torch.bfloat16), InFlexData() + elif dtype == "fp8": + wq = w.to(torch.float8_e4m3fn).transpose(-1, -2).contiguous().transpose(-1, -2) + return ( + wq, + InFlexData(dtype=wq.dtype, scale=w.abs().max().unsqueeze(0)), + MicroscalingCtx(), + ) + else: + assert dtype == "mx4", f"{dtype=}" + swizzle_mx_scale = opt["swizzle_mx_scale"] + swizzle_axis = 2 if swizzle_mx_scale else None + w = w.to(torch.bfloat16) + w, mx_scales, weight_scale_shape = downcast_to_mxfp( + w, torch.uint8, axis=1, swizzle_axis=swizzle_axis + ) + return ( + w, + InFlexData(), + MicroscalingCtx( + weight_scale=mx_scales, + swizzle_mx=swizzle_mx_scale, + actual_weight_scale_shape=weight_scale_shape, + ), + ) + + +def triton_kernel_moe_forward( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_output: TopKOutput, + moe_runner_config: MoeRunnerConfig, + apply_router_weight_on_input: bool = False, + use_fp8_w8a8: bool = False, + per_channel_quant: bool = False, + global_num_experts: int = -1, + expert_map: Optional[torch.Tensor] = None, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[list[int]] = None, +) -> torch.Tensor: + + from sglang.srt.layers.moe.topk import TopKOutputChecker + + assert TopKOutputChecker.format_is_triton_kernel(topk_output) + + routing_data, gather_idx, scatter_idx = topk_output + + return triton_kernel_fused_experts( + hidden_states, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + inplace=False, # triton kernel doesn't support inplace + activation=moe_runner_config.activation, + apply_router_weight_on_input=apply_router_weight_on_input, + use_fp8_w8a8=use_fp8_w8a8, + per_channel_quant=per_channel_quant, + global_num_experts=global_num_experts, + expert_map=expert_map, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + + +# This is a triton implementation of the fused_experts function +def triton_kernel_fused_experts( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + routing_data: RoutingData, + gather_indx: GatherIndx, + scatter_indx: ScatterIndx, + inplace: bool = False, + activation: str = "silu", + apply_router_weight_on_input: bool = False, + use_fp8_w8a8: bool = False, + per_channel_quant: bool = False, + global_num_experts: int = -1, + expert_map: Optional[torch.Tensor] = None, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[list[int]] = None, +) -> torch.Tensor: + + assert use_fp8_w8a8 == False, "use_fp8_w8a8 is not supported" + assert per_channel_quant == False, "per_channel_quant is not supported" + assert expert_map == None, "expert_map is not supported" + assert w1_scale == None, "w1_scale is not supported" + assert w2_scale == None, "w2_scale is not supported" + assert a1_scale == None, "a1_scale is not supported" + assert a2_scale == None, "a2_scale is not supported" + assert block_shape == None, "block_shape is not supported" + + # type check + assert hidden_states.dtype == torch.bfloat16, "hidden_states must be bfloat16" + assert w1.dtype == torch.bfloat16, "w1 must be bfloat16" + assert w2.dtype == torch.bfloat16, "w2 must be bfloat16" + + # Shape check + assert hidden_states.ndim == 2, "hidden_states must be 2D" + assert ( + hidden_states.shape[-1] == w1.shape[-2] + ), f"hidden_states shape[-1] {hidden_states.shape} must be equal to w1 shape[-2] {w1.shape}" + assert ( + w2.shape[-1] == w1.shape[1] + ), f"w2 shape[-1] {w2.shape[-1]} must be equal to w1 shape[1] {w1.shape[1]}" + + # feature check + assert inplace == False, "Inplace is not supported in new triton MoE kernel" + + M, K = hidden_states.shape + E, _, N = w1.shape + n_expts_act = routing_data.n_expts_act + dtype = hidden_states.dtype + + if global_num_experts == -1: + global_num_experts = E + + # consistent with default implementation + intermediate_cache2 = torch.empty( + (M * n_expts_act, N // 2), device="cuda", dtype=dtype + ) + + intermediate_cache1 = matmul_ogs( + hidden_states, + w1, + None, + routing_data, + gather_indx=gather_indx, + gammas=routing_data.gate_scal if apply_router_weight_on_input else None, + ) + + if activation == "silu": + silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2) + elif activation == "gelu": + gelu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2) + else: + raise ValueError(f"Unsupported FusedMoe activation: {activation}") + + intermediate_cache3 = matmul_ogs( + intermediate_cache2, + w2, + None, + routing_data, + scatter_indx=scatter_indx, + gammas=None if apply_router_weight_on_input else routing_data.gate_scal, + ) + + return intermediate_cache3 + + +def triton_kernel_moe_with_bias_forward( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w1_pcg, + b1: torch.Tensor, + w2: torch.Tensor, + w2_pcg, + b2: torch.Tensor, + topk_output: TopKOutput, + moe_runner_config: MoeRunnerConfig, + use_fp8_w8a8: bool = False, + per_channel_quant: bool = False, + global_num_experts: int = -1, + expert_map: Optional[torch.Tensor] = None, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[list[int]] = None, +) -> torch.Tensor: + from sglang.srt.layers.moe.topk import TopKOutputChecker + + assert TopKOutputChecker.format_is_triton_kernel(topk_output) + + routing_data, gather_idx, scatter_idx = topk_output + + return triton_kernel_fused_experts_with_bias( + hidden_states, + w1=w1, + w1_pcg=w1_pcg, + b1=b1, + w2=w2, + w2_pcg=w2_pcg, + b2=b2, + routing_data=routing_data, + gather_indx=gather_idx, + scatter_indx=scatter_idx, + inplace=False, # triton kernel doesn't support inplace + activation=moe_runner_config.activation, + use_fp8_w8a8=use_fp8_w8a8, + per_channel_quant=per_channel_quant, + global_num_experts=global_num_experts, + expert_map=expert_map, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + gemm1_alpha=moe_runner_config.gemm1_alpha, + gemm1_clamp_limit=moe_runner_config.gemm1_clamp_limit, + ) + + +def triton_kernel_fused_experts_with_bias( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w1_pcg, + b1: torch.Tensor, + w2: torch.Tensor, + w2_pcg, + b2: torch.Tensor, + routing_data: RoutingData, + gather_indx: GatherIndx, + scatter_indx: ScatterIndx, + inplace: bool = False, + activation: str = "silu", + use_fp8_w8a8: bool = False, + per_channel_quant: bool = False, + global_num_experts: int = -1, + expert_map: Optional[torch.Tensor] = None, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[list[int]] = None, + gemm1_alpha: Optional[float] = None, + gemm1_clamp_limit: Optional[float] = None, +) -> torch.Tensor: + assert use_fp8_w8a8 == False, "use_fp8_w8a8 is not supported" + assert per_channel_quant == False, "per_channel_quant is not supported" + assert expert_map == None, "expert_map is not supported" + assert w1_scale == None, "w1_scale is not supported" + assert w2_scale == None, "w2_scale is not supported" + assert a1_scale == None, "a1_scale is not supported" + assert a2_scale == None, "a2_scale is not supported" + assert block_shape == None, "block_shape is not supported" + + # type check + assert hidden_states.dtype == torch.bfloat16, "hidden_states must be bfloat16" + for w in (w1, w2): + # TODO assert bf16 or mxfp4 + # assert (w.dtype == torch.bfloat16) or check-is-mxfp4, f"w must be bfloat16 or mxfp4 {w1.dtype=}" + pass + + # Shape check + assert hidden_states.ndim == 2, "hidden_states must be 2D" + assert ( + hidden_states.shape[-1] == w1.shape[-2] + ), f"hidden_states shape[-1] {hidden_states.shape} must be equal to w1 shape[-2] {w1.shape}" + assert ( + w2.shape[-1] == w1.shape[1] + ), f"w2 shape[-1] {w2.shape[-1]} must be equal to w1 shape[1] {w1.shape[1]}" + + # feature check + assert inplace == False, "Inplace is not supported in new triton MoE kernel" + + E, _, _ = w1.shape + + if global_num_experts == -1: + global_num_experts = E + + # TODO maybe completely remove this branch + if w1.dtype == torch.bfloat16: + device = "cuda" + optg = dict() + w1, w1_flex = quantize(w1, "bf16", device, **optg) + w1_pcg = PrecisionConfig(flex_ctx=FlexCtx(rhs_data=w1_flex)) + + w2, w2_flex = quantize(w2, "bf16", device, **optg) + w2_pcg = PrecisionConfig(flex_ctx=FlexCtx(rhs_data=w2_flex)) + + act = FusedActivation( + FnSpecs("swiglu", swiglu_fn, ("alpha", "limit")), + (gemm1_alpha, gemm1_clamp_limit), + 2, + ) + + intermediate_cache = matmul_ogs( + hidden_states, + w1, + b1, + routing_data, + gather_indx=gather_indx, + precision_config=w1_pcg, + gammas=None, + fused_activation=act, + ) + + return matmul_ogs( + intermediate_cache, + w2, + b2, + routing_data, + scatter_indx=scatter_indx, + precision_config=w2_pcg, + gammas=routing_data.gate_scal, + ) diff --git a/python/sglang/srt/layers/moe/moe_runner/__init__.py b/python/sglang/srt/layers/moe/moe_runner/__init__.py new file mode 100644 index 000000000..3320a7875 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/__init__.py @@ -0,0 +1,4 @@ +from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.runner import MoeRunner + +__all__ = ["MoeRunnerConfig", "MoeRunner"] diff --git a/python/sglang/srt/layers/moe/moe_runner/base.py b/python/sglang/srt/layers/moe/moe_runner/base.py new file mode 100644 index 000000000..4d95540e6 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/base.py @@ -0,0 +1,286 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Optional, Tuple, TypeGuard + +import torch + +from sglang.srt.layers.moe.utils import MoeA2ABackend, MoeRunnerBackend + +if TYPE_CHECKING: + from sglang.srt.layers.moe.moe_runner.triton import ( + TritonRunnerCore, + TritonRunnerInput, + TritonRunnerOutput, + ) + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + CombineInputFormat, + DispatchOutput, + DispatchOutputFormat, + ) + + +@dataclass +class MoeRunnerConfig: + + # MoE parameters + num_experts: Optional[int] = None + num_local_experts: Optional[int] = None + hidden_size: Optional[int] = None + intermediate_size_per_partition: Optional[int] = None + layer_id: Optional[int] = None + top_k: Optional[int] = None + num_fused_shared_experts: Optional[int] = None + params_dtype: Optional[torch.dtype] = None + + # Runner configuration + activation: str = "silu" + apply_router_weight_on_input: bool = False + inplace: bool = True + no_combine: bool = False + routed_scaling_factor: Optional[float] = None + gemm1_alpha: Optional[float] = None + gemm1_clamp_limit: Optional[float] = None + + +@dataclass +class RunnerInput(ABC): + + @property + @abstractmethod + def runner_backend(self) -> MoeRunnerBackend: ... + + def runner_backend_is_triton(self) -> TypeGuard[TritonRunnerInput]: + return self.runner_backend == MoeRunnerBackend.TRITON + + +class RunnerOutput(ABC): + + @property + @abstractmethod + def runner_backend(self) -> MoeRunnerBackend: ... + + def runner_backend_is_triton(self) -> TypeGuard[TritonRunnerOutput]: + return self.runner_backend == MoeRunnerBackend.TRITON + + +@dataclass +class MoeQuantInfo(ABC): + """Moe quantization data.""" + + pass + + +class MoeRunnerCore(ABC): + + def __init__(self, config: MoeRunnerConfig): + self.config = config + + @abstractmethod + def run( + self, runner_input: RunnerInput, quant_info: MoeQuantInfo, running_state: dict + ) -> RunnerOutput: + pass + + @property + @abstractmethod + def runner_backend(self) -> MoeRunnerBackend: ... + + def runner_backend_is_triton(self) -> TypeGuard[TritonRunnerCore]: + return self.runner_backend == MoeRunnerBackend.TRITON + + +class FusedOpPool: + + _fused_funcs: dict[str, Callable] = {} + + @classmethod + def register_fused_func( + cls, a2a_backend_name: str, runner_backend_name: str, fused_func: Callable + ): + key = (a2a_backend_name, runner_backend_name) + if key in cls._fused_funcs: + raise ValueError( + f"Fused function for {a2a_backend_name} to {runner_backend_name} is already registered." + ) + assert MoeA2ABackend( + a2a_backend_name + ), f"Invalid dispatch name: {a2a_backend_name}" + assert MoeRunnerBackend( + runner_backend_name + ), f"Invalid runner name: {runner_backend_name}" + cls._fused_funcs[key] = fused_func + + @classmethod + def get_fused_func(cls, dispatch_name: str, runner_name: str) -> Optional[Callable]: + key = (dispatch_name, runner_name) + fused_func = cls._fused_funcs.get(key) + return fused_func + + +class PermuteMethodPool: + + _pre_permute_methods: dict[ + Tuple[DispatchOutputFormat, MoeRunnerBackend], Callable + ] = {} + _post_permute_methods: dict[ + Tuple[MoeRunnerBackend, CombineInputFormat], Callable + ] = {} + + @classmethod + def register_pre_permute( + cls, + dispatch_output_name: str, + runner_backend_name: str, + permute_func: Callable, + ): + """ + Register a customized pre-permute function for the given DispatchOutputFormat and MoeRunnerBackend. + + :param dispatch_output_name: The DispatchOutputFormat name. + :param runner_backend_name: The MoeRunnerBackend name. + :param permute_func: The permute function to register. + """ + # TODO: check if registration is valid + key = (dispatch_output_name, runner_backend_name) + if key in cls._pre_permute_methods: + raise ValueError( + f"Pre-permute method for {dispatch_output_name} to {runner_backend_name} is already registered." + ) + cls._pre_permute_methods[key] = permute_func + + @classmethod + def register_post_permute( + cls, + runner_backend_name: str, + combine_input_name: str, + permute_func: Callable, + ): + """ + Register a customized post-permute function for the given MoeRunnerBackend and CombineInputFormat. + + :param runner_backend_name: The MoeRunnerBackend name. + :param combine_input_name: The CombineInputFormat name. + :param permute_func: The permute function to register. + """ + # TODO: check if registration is valid + key = (runner_backend_name, combine_input_name) + if key in cls._post_permute_methods: + raise ValueError( + f"Post-permute method for {runner_backend_name} to {combine_input_name} is already registered." + ) + cls._post_permute_methods[key] = permute_func + + @classmethod + def get_pre_permute( + cls, + dispatch_output_format: DispatchOutputFormat, + runner_input_format: MoeRunnerBackend, + ) -> Callable: + """ + Retrieve the pre-permute function for the given DispatchOutputFormat and MoeRunnerBackend. + + :param dispatch_output_format: The DispatchOutputFormat type. + :param runner_input_format: The MoeRunnerBackend type. + :return: The registered permute function or None if not found. + """ + key = (dispatch_output_format, runner_input_format) + pre_permute_func = cls._pre_permute_methods.get(key) + assert ( + pre_permute_func is not None + ), f"Pre-permute function for {dispatch_output_format} to {runner_input_format} is not registered" + return pre_permute_func + + @classmethod + def get_post_permute( + cls, + runner_output_format: MoeRunnerBackend, + combine_input_format: CombineInputFormat, + ) -> Callable: + """ + Retrieve the post-permute function for the given MoeRunnerBackend and CombineInputFormat. + + :param runner_output_format: The MoeRunnerBackend type. + :param combine_input_format: The CombineInputFormat type. + :return: The registered permute function or None if not found. + """ + key = (runner_output_format, combine_input_format) + post_permute_func = cls._post_permute_methods.get(key) + assert ( + post_permute_func is not None + ), f"Post-permute function for {runner_output_format} to {combine_input_format} is not registered" + return post_permute_func + + +def register_fused_func( + a2a_backend_name: str, + runner_backend_name: str, +) -> Callable: + """ + Decorator to register a fused function for the given DispatchOutputFormat and MoeRunnerBackend. + + :param a2a_backend_name: The A2A backend name. + :param runner_backend_name: The MoeRunnerBackend name. + :return: The decorator function. + """ + + def decorator(fused_func: Callable): + FusedOpPool.register_fused_func( + a2a_backend_name, runner_backend_name, fused_func + ) + return fused_func + + return decorator + + +def register_pre_permute( + dispatch_output_name: str, + runner_backend_name: str, +) -> Callable: + """ + Decorator to register a pre-permute function for the given DispatchOutputFormat and MoeRunnerBackend. + + :param dispatch_output_name: The DispatchOutputFormat name. + :param runner_backend_name: The MoeRunnerBackend name. + :return: The decorator function. + """ + + def decorator( + permute_func: Callable[ + [DispatchOutput, MoeQuantInfo, MoeRunnerConfig, dict], RunnerInput + ] + ) -> Callable: + + PermuteMethodPool.register_pre_permute( + dispatch_output_name, runner_backend_name, permute_func + ) + return permute_func + + return decorator + + +def register_post_permute( + runner_backend_name: str, + combine_input_name: str, +) -> Callable: + """ + Decorator to register a post-permute function for the given MoeRunnerBackend and CombineInputFormat. + + :param runner_backend_name: The MoeRunnerBackend name. + :param combine_input_name: The CombineInputFormat name. + :return: The decorator function. + """ + + def decorator( + permute_func: Callable[ + [RunnerOutput, MoeQuantInfo, MoeRunnerConfig, dict], CombineInput + ] + ) -> Callable: + PermuteMethodPool.register_post_permute( + runner_backend_name, combine_input_name, permute_func + ) + return permute_func + + return decorator diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py new file mode 100644 index 000000000..3b6fcd980 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import logging +import os +from typing import TYPE_CHECKING + +from sglang.srt.layers.moe.moe_runner.base import ( + FusedOpPool, + MoeRunnerConfig, + PermuteMethodPool, +) +from sglang.srt.layers.moe.moe_runner.triton import TritonRunnerCore +from sglang.srt.layers.moe.utils import get_moe_a2a_backend + +if TYPE_CHECKING: + from sglang.srt.layers.moe.moe_runner.base import MoeQuantInfo + from sglang.srt.layers.moe.token_dispatcher.base import CombineInput, DispatchOutput + from sglang.srt.layers.moe.utils import MoeRunnerBackend + +logger = logging.getLogger(__name__) + + +class MoeRunner: + + def __init__(self, runner_backend: MoeRunnerBackend, config: MoeRunnerConfig): + self.runner_backend = runner_backend + self.config = config + + self.fused_func = None + + if runner_backend.is_triton(): + self.runner_core = TritonRunnerCore(config) + else: + raise NotImplementedError(f"Unsupported runner backend: {runner_backend}") + + a2a_backend_name = get_moe_a2a_backend().value + runner_backend_name = runner_backend.value + + self.fused_func = FusedOpPool.get_fused_func( + a2a_backend_name, runner_backend_name + ) + + SGLANG_CI_DISABLE_MOE_FUSED_FUNC = os.environ.get( + "SGLANG_CI_DISABLE_MOE_FUSED_FUNC", "0" + ) + if SGLANG_CI_DISABLE_MOE_FUSED_FUNC == "1": + logger.info( + "SGLANG_CI_DISABLE_MOE_FUSED_FUNC is set to 1, disabling fused func" + ) + self.fused_func = None + + def run( + self, dispatch_output: DispatchOutput, quant_info: MoeQuantInfo + ) -> CombineInput: + + if self.fused_func is not None: + return self.fused_func(dispatch_output, quant_info, self.config) + + dispatch_format = dispatch_output.format.value + runner_format = self.runner_core.runner_backend.value + self.pre_permute_func = PermuteMethodPool.get_pre_permute( + dispatch_format, runner_format + ) + + running_state = {} + runner_input = self.pre_permute_func( + dispatch_output, quant_info, self.config, running_state + ) + runner_output = self.runner_core.run(runner_input, quant_info, running_state) + + runner_format = self.runner_core.runner_backend.value + combine_format = dispatch_output.format.value + self.post_permute_func = PermuteMethodPool.get_post_permute( + runner_format, combine_format + ) + combine_input = self.post_permute_func( + runner_output, quant_info, self.config, running_state + ) + + return combine_input diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py new file mode 100644 index 000000000..116fdcaa0 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -0,0 +1,448 @@ +from __future__ import annotations + +import functools +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional + +import torch +import triton.language as tl + +from sglang.srt.layers.moe.moe_runner.base import ( + MoeQuantInfo, + MoeRunnerConfig, + MoeRunnerCore, + RunnerInput, + RunnerOutput, + register_fused_func, + register_post_permute, + register_pre_permute, +) +from sglang.srt.layers.moe.utils import MoeRunnerBackend +from sglang.srt.utils import cpu_has_amx_support, is_cpu, is_cuda, is_hip + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher.standard import ( + StandardCombineInput, + StandardDispatchOutput, + ) + + +_is_hip = is_hip() +_is_cuda = is_cuda() +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() +_use_aiter = bool(int(os.getenv("SGLANG_MOE_USE_AITER", "0"))) +_MOE_PADDING_SIZE = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0 + + +if _is_cuda: + from sgl_kernel import gelu_and_mul, silu_and_mul +elif _is_cpu and _is_cpu_amx_available: + pass +elif _is_hip: + from vllm import _custom_ops as vllm_ops # gelu_and_mul, silu_and_mul + + if _use_aiter: + try: + from aiter import moe_sum + except ImportError: + raise ImportError("aiter is required when SGLANG_USE_AITER is set to True") + + +if _is_cuda or _is_hip: + from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size + + +@dataclass +class TritonRunnerInput(RunnerInput): + + hidden_states: torch.Tensor + topk_weights: torch.Tensor + topk_ids: torch.Tensor + sorted_token_ids: torch.Tensor + expert_ids: torch.Tensor + num_tokens_post_padded: torch.Tensor + + @property + def runner_backend(self) -> MoeRunnerBackend: + return MoeRunnerBackend.TRITON + + +@dataclass +class TritonRunnerOutput(RunnerOutput): + + hidden_states: torch.Tensor + + @property + def runner_backend(self) -> MoeRunnerBackend: + return MoeRunnerBackend.TRITON + + +@dataclass +class TritonMoeQuantInfo(MoeQuantInfo): + w13_weight: torch.Tensor + w2_weight: torch.Tensor + b13: Optional[torch.Tensor] = None + b2: Optional[torch.Tensor] = None + use_fp8_w8a8: bool = False + use_int8_w8a8: bool = False + use_int8_w8a16: bool = False + use_int4_w4a16: bool = False + per_channel_quant: bool = False + w13_scale: Optional[torch.Tensor] = None + w2_scale: Optional[torch.Tensor] = None + w13_zp: Optional[torch.Tensor] = None + w2_zp: Optional[torch.Tensor] = None + a13_scale: Optional[torch.Tensor] = None + a2_scale: Optional[torch.Tensor] = None + block_shape: Optional[List[int]] = None + + +class TritonRunnerCore(MoeRunnerCore): + + def __init__(self, config: MoeRunnerConfig): + super().__init__(config) + + def run( + self, + runner_input: TritonRunnerInput, + quant_info: TritonMoeQuantInfo, + running_state: dict, + ) -> TritonRunnerOutput: + + # TODO: move these functions to the triton runner + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + invoke_fused_moe_kernel, + moe_sum_reduce_torch_compile, + moe_sum_reduce_triton, + swiglu_with_alpha_and_limit, + ) + + hidden_states = runner_input.hidden_states + topk_weights = runner_input.topk_weights + topk_ids = runner_input.topk_ids + sorted_token_ids = runner_input.sorted_token_ids + expert_ids = runner_input.expert_ids + num_tokens_post_padded = runner_input.num_tokens_post_padded + + w13 = quant_info.w13_weight + w2 = quant_info.w2_weight + b13 = quant_info.b13 + b2 = quant_info.b2 + a13_scale = quant_info.a13_scale + a2_scale = quant_info.a2_scale + w13_scale = quant_info.w13_scale + w2_scale = quant_info.w2_scale + w13_zp = quant_info.w13_zp + w2_zp = quant_info.w2_zp + block_shape = quant_info.block_shape + per_channel_quant = quant_info.per_channel_quant + use_fp8_w8a8 = quant_info.use_fp8_w8a8 + use_int8_w8a8 = quant_info.use_int8_w8a8 + use_int8_w8a16 = quant_info.use_int8_w8a16 + use_int4_w4a16 = quant_info.use_int4_w4a16 + + activation = self.config.activation + no_combine = self.config.no_combine + inplace = self.config.inplace + gemm1_alpha = self.config.gemm1_alpha + gemm1_limit = self.config.gemm1_clamp_limit + routed_scaling_factor = self.config.routed_scaling_factor + apply_router_weight_on_input = self.config.apply_router_weight_on_input + + M = hidden_states.shape[0] + E, N, _ = w13.shape + compute_type = ( + tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16 + ) + + intermediate_cache1 = torch.empty( + (M, topk_ids.shape[1], N), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + invoke_fused_moe_kernel( + hidden_states, + w13, + b13, + intermediate_cache1, + a13_scale, + w13_scale, + w13_zp, + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + apply_router_weight_on_input, + topk_ids.shape[1], + running_state["config"], + compute_type=compute_type, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + use_int4_w4a16=use_int4_w4a16, + per_channel_quant=per_channel_quant, + block_shape=block_shape, + ) + + intermediate_cache2 = torch.empty( + (M * topk_ids.shape[1], N // 2), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + if activation == "silu": + if gemm1_alpha is not None: + assert gemm1_limit is not None + intermediate_cache2 = swiglu_with_alpha_and_limit( + intermediate_cache1.view(-1, N), + gemm1_alpha, + gemm1_limit, + ) + elif _is_cuda: + silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2) + else: + vllm_ops.silu_and_mul( + intermediate_cache2, intermediate_cache1.view(-1, N) + ) + elif activation == "gelu": + assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu" + assert gemm1_limit is None, "gemm1_limit is not supported for gelu" + if _is_cuda: + gelu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2) + else: + vllm_ops.gelu_and_mul( + intermediate_cache2, intermediate_cache1.view(-1, N) + ) + else: + raise ValueError(f"Unsupported activation: {activation=}") + + intermediate_cache3 = torch.empty( + (M, topk_ids.shape[1], w2.shape[1]), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + if no_combine: + assert not inplace + out_hidden_states = torch.empty( + (M, topk_ids.shape[1], w2.shape[1]), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + elif inplace: + out_hidden_states = hidden_states + else: + out_hidden_states = torch.empty_like(hidden_states) + + invoke_fused_moe_kernel( + intermediate_cache2, + w2, + b2, + ( + intermediate_cache3 + if not no_combine and topk_ids.shape[1] != 1 + else out_hidden_states.unsqueeze(0) + ), + a2_scale, + w2_scale, + w2_zp, + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + not apply_router_weight_on_input, + 1, + running_state["config"], + compute_type=compute_type, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + use_int4_w4a16=use_int4_w4a16, + per_channel_quant=per_channel_quant, + block_shape=block_shape, + ) + + if routed_scaling_factor is None: + routed_scaling_factor = 1.0 + + if no_combine: + pass + elif _is_cuda: + if topk_ids.shape[1] == 1 and routed_scaling_factor == 1.0: + pass # we write directly into out_hidden_states + elif topk_ids.shape[1] == 2 and routed_scaling_factor == 1.0: + torch.add( + intermediate_cache3[:, 0], + intermediate_cache3[:, 1], + out=out_hidden_states, + ).squeeze(dim=1) + else: + # According to micro benchmark results, torch.compile can get better performance for small token. + if M <= 32: + moe_sum_reduce_torch_compile( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states, + routed_scaling_factor, + ) + else: + moe_sum_reduce_triton( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states, + routed_scaling_factor, + ) + elif _is_hip: + if _use_aiter: + moe_sum( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states, + ) + else: + vllm_ops.moe_sum( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states, + ) + else: + vllm_ops.moe_sum( + intermediate_cache3.view(*intermediate_cache3.shape), + out_hidden_states, + ) + + return TritonRunnerOutput( + hidden_states=out_hidden_states, + ) + + @property + def runner_backend(self) -> MoeRunnerBackend: + return MoeRunnerBackend.TRITON + + +@register_fused_func("none", "triton") +def fused_experts_none_to_triton( + dispatch_output: StandardDispatchOutput, + quant_info: TritonMoeQuantInfo, + runner_config: MoeRunnerConfig, +) -> StandardCombineInput: + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_experts + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + + output = fused_experts( + hidden_states=dispatch_output.hidden_states, + w1=quant_info.w13_weight, + w2=quant_info.w2_weight, + topk_output=dispatch_output.topk_output, + moe_runner_config=runner_config, + b1=quant_info.b13, + b2=quant_info.b2, + use_fp8_w8a8=quant_info.use_fp8_w8a8, + use_int8_w8a8=quant_info.use_int8_w8a8, + use_int8_w8a16=quant_info.use_int8_w8a16, + use_int4_w4a16=quant_info.use_int4_w4a16, + per_channel_quant=quant_info.per_channel_quant, + w1_scale=quant_info.w13_scale, + w2_scale=quant_info.w2_scale, + w1_zp=quant_info.w13_zp, + w2_zp=quant_info.w2_zp, + a1_scale=quant_info.a13_scale, + a2_scale=quant_info.a2_scale, + block_shape=quant_info.block_shape, + ) + + return StandardCombineInput( + hidden_states=output, + ) + + +@register_pre_permute("standard", "triton") +def pre_permute_standard_to_triton( + dispatch_output: StandardDispatchOutput, + quant_info: TritonMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> TritonRunnerInput: + + # NOTE: this is dead code as a fused func for standard format is registered. + # This is left here for testing and examples. + + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + get_config_dtype_str, + moe_align_block_size, + try_get_optimal_moe_config, + ) + from sglang.srt.layers.moe.topk import TopKOutputChecker + + hidden_states, topk_output = dispatch_output + + assert TopKOutputChecker.format_is_standard(topk_output) + + num_tokens = hidden_states.shape[0] + num_local_experts = runner_config.num_local_experts + + if ( + not (quant_info.use_fp8_w8a8 or quant_info.use_int8_w8a8) + or quant_info.block_shape is not None + or _use_aiter + ): + padding_size = 0 + else: + padding_size = _MOE_PADDING_SIZE + + config_dtype = get_config_dtype_str( + use_fp8_w8a8=quant_info.use_fp8_w8a8, + use_int8_w8a8=quant_info.use_int8_w8a8, + use_int8_w8a16=quant_info.use_int8_w8a16, + use_int4_w4a16=quant_info.use_int4_w4a16, + dtype=hidden_states.dtype, + ) + + get_config_func = functools.partial( + try_get_optimal_moe_config, + quant_info.w13_weight.shape, + ( + num_local_experts, + quant_info.w2_weight.shape[1], + quant_info.w2_weight.shape[2] - padding_size, + ), + topk_output.topk_ids.shape[1], + config_dtype, + block_shape=quant_info.block_shape, + ) + + config = get_config_func(num_tokens) + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_output.topk_ids, config["BLOCK_SIZE_M"], num_local_experts + ) + + running_state["config"] = config + + return TritonRunnerInput( + hidden_states=hidden_states, + topk_weights=topk_output.topk_weights, + topk_ids=topk_output.topk_ids, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + ) + + +@register_post_permute("triton", "standard") +def post_permute_triton_to_standard( + runner_output: TritonRunnerOutput, + quant_info: TritonMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> StandardCombineInput: + + # NOTE: this is dead code as a fused func for standard format is registered. + # This is left here for testing and examples. + + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + + return StandardCombineInput( + hidden_states=runner_output.hidden_states, + ) diff --git a/python/sglang/srt/layers/moe/rocm_moe_utils.py b/python/sglang/srt/layers/moe/rocm_moe_utils.py new file mode 100644 index 000000000..5fe2de1e5 --- /dev/null +++ b/python/sglang/srt/layers/moe/rocm_moe_utils.py @@ -0,0 +1,141 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.9.1rc2/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from enum import IntEnum +from functools import cache +from typing import Optional + +import torch + +from sglang.srt.utils import direct_register_custom_op, get_bool_env_var, is_hip + +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + + +class ActivationMethod(IntEnum): + # This allows interfacing with AITER ActivationType enum + # without importing the ActivationType enum from AITER globally. + SILU = 0 + GELU = 1 + + +def rocm_aiter_asm_moe_tkw1_impl( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + fc1_scale: Optional[torch.Tensor] = None, + fc2_scale: Optional[torch.Tensor] = None, + fc1_smooth_scale: Optional[torch.Tensor] = None, + fc2_smooth_scale: Optional[torch.Tensor] = None, + a16: bool = False, + per_tensor_quant_scale: Optional[torch.Tensor] = None, + expert_mask: Optional[torch.Tensor] = None, + activation_method: int = ActivationMethod.SILU.value, +) -> torch.Tensor: + + from aiter import ActivationType + from aiter.fused_moe_bf16_asm import asm_moe_tkw1 + + activation = ActivationType(activation_method) + + return asm_moe_tkw1( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + fc1_scale=fc1_scale, + fc2_scale=fc2_scale, + fc1_smooth_scale=fc1_smooth_scale, + fc2_smooth_scale=fc2_smooth_scale, + a16=a16, + per_tensor_quant_scale=per_tensor_quant_scale, + expert_mask=expert_mask, + activation=activation, + ) + + +def rocm_aiter_asm_moe_tkw1_fake( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + fc1_scale: Optional[torch.Tensor] = None, + fc2_scale: Optional[torch.Tensor] = None, + fc1_smooth_scale: Optional[torch.Tensor] = None, + fc2_smooth_scale: Optional[torch.Tensor] = None, + a16: bool = False, + per_tensor_quant_scale: Optional[torch.Tensor] = None, + expert_mask: Optional[torch.Tensor] = None, + activation_method: int = ActivationMethod.SILU.value, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +if _use_aiter: + + direct_register_custom_op( + op_name="rocm_aiter_asm_moe_tkw1", + op_func=rocm_aiter_asm_moe_tkw1_impl, + mutates_args=[], + fake_impl=rocm_aiter_asm_moe_tkw1_fake, + ) + + +def rocm_fused_experts_tkw1( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str = "silu", + apply_router_weight_on_input: bool = False, + use_fp8_w8a8: bool = False, + per_channel_quant: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, + block_shape: Optional[list[int]] = None, +) -> torch.Tensor: + + activation_method = ( + ActivationMethod.SILU if activation == "silu" else ActivationMethod.GELU + ) + # All AITER Fused MoE kernels are expecting the following datatypes + topk_weights = topk_weights.to(torch.float32) + topk_ids = topk_ids.to(torch.int32) + + # w8a8 per-channel quantization + if per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8: + # AITER tkw1 kernel for FP8 models with `apply_router_weight_on_input` + # This applies topk_weights on the GEMM output of the first FC layer + # rather than the second FC. + assert ( + topk_weights.dim() == 2 + ), "`topk_weights` should be in shape (num_tokens, topk)" + assert topk_weights.shape[-1] == 1, ( + "Only support topk=1 when" " `apply_router_weight_on_input` is True" + ) + + return torch.ops.sglang.rocm_aiter_asm_moe_tkw1( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + fc1_scale=w1_scale, + fc2_scale=w2_scale, + fc1_smooth_scale=None, + fc2_smooth_scale=None, + a16=False, + per_tensor_quant_scale=None, + expert_mask=None, + activation_method=activation_method, + ) + else: + assert False, "This should not be called." diff --git a/python/sglang/srt/layers/moe/router.py b/python/sglang/srt/layers/moe/router.py new file mode 100644 index 000000000..0138dcdad --- /dev/null +++ b/python/sglang/srt/layers/moe/router.py @@ -0,0 +1,392 @@ +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.moe.topk import fused_topk +from sglang.srt.utils import is_hip + +_is_hip = is_hip() + + +@triton.jit +def fused_moe_router_kernel( + input_ptr, # input (bs, hidden_dim) + moe_router_weight_ptr, # input (num_experts, hidden_dim) + topk_weights_ptr, # output (bs, topk) + topk_ids_ptr, # output (bs, topk) + correction_bias_ptr, + is_correction_bias: tl.constexpr, + num_experts: tl.constexpr, + topk: tl.constexpr, + moe_softcapping: tl.constexpr, + moe_renormalize: tl.constexpr, # not supported + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_dim + + # moe_router_weight is k major + expert_offsets = tl.arange(0, num_experts)[:, None] + router_mask = mask[None, :] + w_router = tl.load( + moe_router_weight_ptr + expert_offsets * hidden_dim + offsets[None, :], + mask=router_mask, + other=0.0, + ) + + x = tl.load(input_ptr + pid * hidden_dim + offsets, mask=mask, other=0.0) + + # todo: tl.dot? + logits = tl.sum((w_router.to(tl.float32) * x[None, :].to(tl.float32)), axis=-1) + + # logit softcap + if moe_softcapping == 0: + logits_softcapped = logits + else: + logits_scaled = logits / moe_softcapping + exped = tl.exp(2 * logits_scaled) + top = exped - 1 + bottom = exped + 1 + logits_softcapped = top / bottom * moe_softcapping + + # Add bias after softcapping + if is_correction_bias: + bias = tl.load(correction_bias_ptr + tl.arange(0, num_experts)) + logits_softcapped = logits_softcapped + bias + + # topk + # assert 1 <= topk <= num_experts + + # 5.38 us + + top1 = tl.argmax(logits_softcapped, axis=0) + tl.store(topk_ids_ptr + pid * topk + 0, top1) # 5.63 us + + top1_v = tl.max(logits_softcapped, axis=0) + invsumexp = 1.0 / tl.sum(tl.exp(logits_softcapped - top1_v), axis=0) + + tl.store( + topk_weights_ptr + pid * topk + 0, + invsumexp, + ) # 5.73 us + + if topk >= 2: + top2 = tl.argmax( + tl.where( + tl.arange(0, num_experts) != top1, logits_softcapped, float("-inf") + ), + axis=0, + ) + tl.store(topk_ids_ptr + pid * topk + 1, top2) + top2_v = tl.sum(logits_softcapped * (tl.arange(0, num_experts) == top2), axis=0) + tl.store( + topk_weights_ptr + pid * topk + 1, + tl.exp(top2_v - top1_v) * invsumexp, + ) # 5.95us + + # probably slow + if topk > 2: + topk_mask = tl.full(logits_softcapped.shape, 1.0, dtype=logits_softcapped.dtype) + topk_mask = tl.where( + tl.arange(0, num_experts) != top1, topk_mask, float("-inf") + ) + topk_mask = tl.where( + tl.arange(0, num_experts) != top2, topk_mask, float("-inf") + ) + for i in range(2, topk): + topi = tl.argmax(logits_softcapped + topk_mask, axis=0) + topk_mask = tl.where( + tl.arange(0, num_experts) != topi, topk_mask, float("-inf") + ) + tl.store(topk_ids_ptr + pid * topk + i, topi) + topi_v = tl.sum( + logits_softcapped * (tl.arange(0, num_experts) == topi), axis=0 + ) + tl.store( + topk_weights_ptr + pid * topk + i, + tl.exp(topi_v - top1_v) * invsumexp, + ) + # assert not moe_renormalize, "moe weight renormalization not implemented" + + +def fused_moe_router_impl( + x: torch.Tensor, + router_weight: torch.Tensor, + topk: int, + moe_softcapping: float, + correction_bias: Optional[torch.Tensor] = None, +): + assert len(x.shape) == 2 and x.shape[1] == router_weight.shape[1] + bs, hidden_dim = x.shape + num_experts = router_weight.shape[0] + + # router_logits = torch.empty((bs, num_experts), dtype=torch.float32, device=x.device) + topk_weights = torch.empty((bs, topk), dtype=torch.float32, device=x.device) + topk_ids = torch.empty((bs, topk), dtype=torch.int32, device=x.device) + is_correction_bias = correction_bias is not None + + max_warps = 16 if _is_hip else 32 + config = { + "BLOCK_SIZE": triton.next_power_of_2(hidden_dim), + "num_warps": max( + min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), max_warps), 4 + ), + } + + fused_moe_router_kernel[(bs,)]( + x, + router_weight, + topk_weights, + topk_ids, + correction_bias, + is_correction_bias=is_correction_bias, + num_experts=num_experts, + topk=topk, + moe_softcapping=moe_softcapping, + moe_renormalize=False, + hidden_dim=hidden_dim, + **config, + ) + + return topk_weights, topk_ids + + +@triton.jit +def fused_moe_router_large_bs_kernel( + a_ptr, # input (bs, hidden_dim) + b_ptr, # input (num_experts, hidden_dim) + topk_weights_ptr, # output (bs, topk) + topk_ids_ptr, # output (bs, topk) + bs, + num_experts: tl.constexpr, + topk: tl.constexpr, # only support topk <= 2 + moe_softcapping: tl.constexpr, + moe_renormalize: tl.constexpr, # not supported + K: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + stride_am: tl.constexpr, + stride_bn: tl.constexpr, +): + + # 1. get block id + pid = tl.program_id(axis=0) + + # 2. create pointers for the first block of A and B + # 2.1. setup a_ptrs with offsets in m and k + offs_m = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)[:, None] + bs_mask = offs_m < bs + offs_k = tl.arange(0, BLOCK_SIZE_K)[None, :] + a_ptrs = a_ptr + (offs_m * stride_am + offs_k) + + # 2.2. setup b_ptrs with offsets in k and n. + # Note: b matrix is k-major. + offs_k = tl.arange(0, BLOCK_SIZE_K)[None, :] + offs_n = tl.arange(0, BLOCK_SIZE_N)[:, None] + expert_mask = offs_n < num_experts + b_ptrs = b_ptr + (offs_n * stride_bn + offs_k) + + # 3. Create an accumulator of float32 of size [BLOCK_SIZE_M, BLOCK_SIZE_N] + # 3.1. iterate in K dimension + # 3.2. transpose tile B + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, K // BLOCK_SIZE_K): # hidden_dim % BLOCK_SIZE_K == 0 + a = tl.load( + a_ptrs, + mask=bs_mask, + other=0.0, + ).to(tl.float32) + b = tl.load(b_ptrs, mask=expert_mask, other=0.0).to(tl.float32).T + acc += tl.dot(a, b) + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K + b_ptrs += BLOCK_SIZE_K + + # 4. logit softcap + if moe_softcapping == 0: + logits_softcapped = acc + else: + logits_scaled = acc / moe_softcapping + exped = tl.exp(2 * logits_scaled) + logits_softcapped = (exped - 1) / (exped + 1) * moe_softcapping + + # 5. top1 + arange_block_size_n = tl.arange(0, BLOCK_SIZE_N)[None, :] + cond_top1 = arange_block_size_n < num_experts + top1 = tl.argmax(tl.where(cond_top1, logits_softcapped, float("-inf")), axis=1) + top1_v = tl.max( + tl.where(cond_top1, logits_softcapped, float("-inf")), axis=1, keep_dims=True + ) + top1_invsumexp = 1.0 / tl.sum( + tl.where(cond_top1, tl.exp(logits_softcapped - top1_v), 0.0), axis=1 + ) + + # 6. store top1 to output + offs_top1 = pid * topk * BLOCK_SIZE_M + topk * tl.arange(0, BLOCK_SIZE_M) + top1_mask = offs_top1 < bs * topk + tl.store(topk_ids_ptr + offs_top1, top1, mask=top1_mask) + tl.store( + topk_weights_ptr + offs_top1, + top1_invsumexp, + mask=top1_mask, + ) + + # 7. handle topk == 2 + if topk == 2: + cond_top2 = (arange_block_size_n < num_experts) & ( + arange_block_size_n != top1[:, None] + ) + top2 = tl.argmax( + tl.where(cond_top2, logits_softcapped, float("-inf")), + axis=1, + keep_dims=True, + ) + top2_v = tl.sum( + logits_softcapped * (arange_block_size_n == top2), axis=1, keep_dims=True + ) + top2_invsumexp = tl.exp(top2_v - top1_v) * top1_invsumexp[:, None] + + # store top2 + offs_top2 = ( + pid * topk * BLOCK_SIZE_M + topk * tl.arange(0, BLOCK_SIZE_M)[:, None] + 1 + ) + top2_mask = offs_top2 < bs * topk + tl.store(topk_ids_ptr + offs_top2, top2, mask=top2_mask) + tl.store( + topk_weights_ptr + offs_top2, + top2_invsumexp, + mask=top2_mask, + ) + + +def fused_moe_router_large_bs_impl( + x: torch.Tensor, + router_weight: torch.Tensor, + topk: int, + moe_softcapping: float, + BLOCK_SIZE_M: int, + BLOCK_SIZE_N: int, + BLOCK_SIZE_K: int, +): + assert len(x.shape) == 2 and x.shape[1] == router_weight.shape[1] + bs, hidden_dim = x.shape + num_experts = router_weight.shape[0] + + assert num_experts <= BLOCK_SIZE_N + assert hidden_dim % BLOCK_SIZE_K == 0 + assert topk <= 2 + + topk_weights = torch.empty((bs, topk), dtype=torch.float32, device=x.device) + topk_ids = torch.empty((bs, topk), dtype=torch.int32, device=x.device) + + grid = (triton.cdiv(bs, BLOCK_SIZE_M) * triton.cdiv(num_experts, BLOCK_SIZE_N),) + + fused_moe_router_large_bs_kernel[grid]( + a_ptr=x, + b_ptr=router_weight, + topk_weights_ptr=topk_weights, + topk_ids_ptr=topk_ids, + bs=bs, + num_experts=num_experts, + topk=topk, + moe_softcapping=moe_softcapping, + moe_renormalize=False, + K=hidden_dim, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + stride_am=hidden_dim, + stride_bn=hidden_dim, + ) + + return topk_weights, topk_ids + + +def fused_moe_router_shim( + moe_softcapping, + hidden_states, + gating_output, + topk, + renormalize, + correction_bias: Optional[torch.Tensor] = None, +): + assert not renormalize + assert ( + len(hidden_states.shape) == 2 + and hidden_states.shape[1] == gating_output.shape[1] + ) + bs, hidden_dim = hidden_states.shape + num_experts = gating_output.shape[0] + BLOCK_SIZE_M = 32 + BLOCK_SIZE_N = 16 + BLOCK_SIZE_K = 256 + if ( + bs >= 512 + and topk <= 2 + and num_experts <= BLOCK_SIZE_N + and hidden_dim % BLOCK_SIZE_K == 0 + ): + return fused_moe_router_large_bs_impl( + x=hidden_states, + router_weight=gating_output, + topk=topk, + moe_softcapping=moe_softcapping, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + ) + else: + return fused_moe_router_impl( + x=hidden_states, + router_weight=gating_output, + topk=topk, + moe_softcapping=moe_softcapping, + correction_bias=correction_bias, + ) + + +class FusedMoeRouter: + def __init__(self, router_linear, topk, moe_softcapping) -> None: + self.router_linear = router_linear + self.topk = topk + self.moe_softcapping = moe_softcapping + + def __call__(self, *args, **kwargs): + return self.forward(*args, **kwargs) + + def forward( + self, x: torch.Tensor, residual: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + if x.is_cuda: + return self.forward_cuda(x, residual) + else: + return self.forward_vllm(x, residual) + + def forward_cuda( + self, x: torch.Tensor, autotune=False + ) -> Tuple[torch.Tensor, torch.Tensor]: + return fused_moe_router_shim( + moe_softcapping=self.moe_softcapping, + hidden_states=x, + gating_output=self.router_linear.weight, + topk=self.topk, + renormalize=False, + ) + + def forward_vllm( + self, + x: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # g, _ = self.router_linear.forward(x) + g = x.float() @ self.router_linear.weight.T.float() + + g = torch.tanh(g.float() / self.moe_softcapping) * self.moe_softcapping + + return fused_topk(x, g, self.topk, False) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py new file mode 100644 index 000000000..e1dbcdd44 --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py @@ -0,0 +1,41 @@ +from sglang.srt.layers.moe.token_dispatcher.base import ( + BaseDispatcher, + BaseDispatcherConfig, + CombineInput, + CombineInputChecker, + CombineInputFormat, + DispatchOutput, + DispatchOutputChecker, + DispatchOutputFormat, +) +from sglang.srt.layers.moe.token_dispatcher.deepep import ( + DeepEPConfig, + DeepEPDispatcher, + DeepEPLLCombineInput, + DeepEPLLOutput, + DeepEPNormalCombineInput, + DeepEPNormalOutput, +) +from sglang.srt.layers.moe.token_dispatcher.standard import ( + StandardCombineInput, + StandardDispatchOutput, +) + +__all__ = [ + "BaseDispatcher", + "BaseDispatcherConfig", + "CombineInput", + "CombineInputChecker", + "CombineInputFormat", + "DispatchOutput", + "DispatchOutputFormat", + "DispatchOutputChecker", + "StandardDispatchOutput", + "StandardCombineInput", + "DeepEPConfig", + "DeepEPDispatcher", + "DeepEPNormalOutput", + "DeepEPLLOutput", + "DeepEPLLCombineInput", + "DeepEPNormalCombineInput", +] diff --git a/python/sglang/srt/layers/moe/token_dispatcher/base.py b/python/sglang/srt/layers/moe/token_dispatcher/base.py new file mode 100644 index 000000000..155860886 --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/base.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from enum import Enum +from typing import TYPE_CHECKING, Protocol, TypeGuard, Union, runtime_checkable + +import torch + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + DeepEPLLCombineInput, + DeepEPLLOutput, + DeepEPNormalCombineInput, + DeepEPNormalOutput, + StandardCombineInput, + StandardDispatchOutput, + ) + from sglang.srt.layers.moe.topk import TopKOutput + +# ------------------------------ Dispatch Output ------------------------------------- + + +class DispatchOutputChecker: + + @staticmethod + def format_is_standard( + dispatch_output: DispatchOutput, + ) -> TypeGuard[StandardDispatchOutput]: + return dispatch_output.format.is_standard() + + @staticmethod + def format_is_deepep_normal( + dispatch_output: DispatchOutput, + ) -> TypeGuard[DeepEPNormalOutput]: + return dispatch_output.format.is_deepep_normal() + + @staticmethod + def format_is_deepep_ll( + dispatch_output: DispatchOutput, + ) -> TypeGuard[DeepEPLLOutput]: + return dispatch_output.format.is_deepep_ll() + + @staticmethod + def format_is_deepep( + dispatch_output: DispatchOutput, + ) -> TypeGuard[Union[DeepEPNormalOutput, DeepEPLLOutput]]: + return dispatch_output.format.is_deepep() + + +class DispatchOutputFormat(Enum): + + STANDARD = "standard" + DEEPEP_NORMAL = "deepep_normal" + DEEPEP_LL = "deepep_ll" + + def is_standard(self) -> bool: + return self == DispatchOutputFormat.STANDARD + + def is_deepep_normal(self) -> bool: + return self == DispatchOutputFormat.DEEPEP_NORMAL + + def is_deepep_ll(self) -> bool: + return self == DispatchOutputFormat.DEEPEP_LL + + def is_deepep(self) -> bool: + return self in [ + DispatchOutputFormat.DEEPEP_NORMAL, + DispatchOutputFormat.DEEPEP_LL, + ] + + +@runtime_checkable +class DispatchOutput(Protocol): + """Protocol for dispatch outputs in different formats.""" + + # TODO: add hidden_states to the protocol + + @property + def format(self) -> DispatchOutputFormat: ... + + +# ------------------------------ Combine Input ------------------------------------- + + +class CombineInputChecker: + @staticmethod + def format_is_standard( + combine_input: CombineInput, + ) -> TypeGuard[StandardCombineInput]: + return combine_input.format == CombineInputFormat.STANDARD + + @staticmethod + def format_is_deepep_normal( + combine_input: CombineInput, + ) -> TypeGuard[DeepEPNormalCombineInput]: + return combine_input.format == CombineInputFormat.DEEPEP_NORMAL + + @staticmethod + def format_is_deepep_ll( + combine_input: CombineInput, + ) -> TypeGuard[DeepEPLLCombineInput]: + return combine_input.format == CombineInputFormat.DEEPEP_LL + + @staticmethod + def format_is_deepep( + combine_input: CombineInput, + ) -> TypeGuard[Union[DeepEPNormalCombineInput, DeepEPLLCombineInput]]: + return combine_input.format in [ + CombineInputFormat.DEEPEP_NORMAL, + CombineInputFormat.DEEPEP_LL, + ] + + +class CombineInputFormat(Enum): + STANDARD = "standard" + DEEPEP_NORMAL = "deepep_normal" + DEEPEP_LL = "deepep_ll" + + +@runtime_checkable +class CombineInput(Protocol): + """Protocol for combine inputs in different formats.""" + + # TODO: add hidden_states to the protocol + + @property + def format(self) -> CombineInputFormat: ... + + +# ------------------------------ Base Dispatcher ------------------------------------- + + +class BaseDispatcherConfig(ABC): + """Base class for dispatcher configs.""" + + pass + + +class BaseDispatcher(ABC): + """Base class for dispatchers.""" + + @abstractmethod + def dispatch( + self, hidden_states: torch.Tensor, topk_output: TopKOutput, **kwargs + ) -> DispatchOutput: + pass + + @abstractmethod + def combine(self, combine_input: CombineInput, **kwargs) -> torch.Tensor: + pass diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py new file mode 100644 index 000000000..450cff0cb --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -0,0 +1,732 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union + +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers.moe.token_dispatcher.base import ( + BaseDispatcher, + BaseDispatcherConfig, + CombineInput, + CombineInputFormat, + DispatchOutput, + DispatchOutputFormat, +) +from sglang.srt.layers.moe.utils import DeepEPMode, get_deepep_config, is_tbo_enabled +from sglang.srt.layers.quantization import deep_gemm_wrapper +from sglang.srt.utils import ( + get_bool_env_var, + get_int_env_var, + is_hip, + is_npu, + load_json_config, +) + +_is_npu = is_npu() + +try: + from deep_ep import Buffer, Config + + if not _is_npu: + from sglang.srt.layers.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, + ) + + use_deepep = True +except ImportError: + use_deepep = False + +from enum import Enum, IntEnum, auto + +import torch +import torch.distributed as dist + +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() + +logger = logging.getLogger(__name__) + + +class DeepEPNormalOutput(NamedTuple): + """DeepEP normal dispatch output.""" + + hidden_states: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] + # hidden_states_scale + topk_idx: torch.Tensor + topk_weights: torch.Tensor + num_recv_tokens_per_expert: List[int] + + @property + def format(self) -> DispatchOutputFormat: + return DispatchOutputFormat.DEEPEP_NORMAL + + +class DeepEPLLOutput(NamedTuple): + """DeepEP low latency dispatch output.""" + + hidden_states_fp8: Tuple[torch.Tensor, torch.Tensor] + topk_idx: torch.Tensor + topk_weights: torch.Tensor + masked_m: torch.Tensor + expected_m: int + + @property + def format(self) -> DispatchOutputFormat: + return DispatchOutputFormat.DEEPEP_LL + + +assert isinstance(DeepEPNormalOutput, DispatchOutput) +assert isinstance(DeepEPLLOutput, DispatchOutput) + + +class DeepEPNormalCombineInput(NamedTuple): + """DeepEP normal combine input.""" + + pass + + @property + def format(self) -> CombineInputFormat: + return CombineInputFormat.DEEPEP_NORMAL + + +class DeepEPLLCombineInput(NamedTuple): + """DeepEP low latency combine input.""" + + pass + + @property + def format(self) -> CombineInputFormat: + return CombineInputFormat.DEEPEP_LL + + +assert isinstance(DeepEPNormalCombineInput, CombineInput) +assert isinstance(DeepEPLLCombineInput, CombineInput) + + +class DeepEPDispatchMode(IntEnum): + NORMAL = auto() + LOW_LATENCY = auto() + + +class DeepEPBuffer: + _buffer = None + _dispatch_mode: Optional[DeepEPDispatchMode] = None + _hidden_size: Optional[int] = None + _num_max_dispatch_tokens_per_rank: Optional[int] = None + _num_experts: Optional[int] = None + + @classmethod + def get_deepep_buffer( + cls, + group: dist.ProcessGroup, + hidden_size: int, + param_bytes: int, + deepep_mode: DeepEPMode, + num_max_dispatch_tokens_per_rank: int = -1, + num_experts: int = -1, + ): + if cls._buffer is not None: + return cls._buffer + + cls._hidden_size = hidden_size + cls._num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank + cls._num_experts = num_experts + + num_nvl_bytes, num_rdma_bytes = 0, 0 + if deepep_mode.enable_normal(): + hidden_bytes = hidden_size * param_bytes + for config in ( + DeepEPConfig.get_instance().normal_dispatch_config + or Buffer.get_dispatch_config(group.size()), + DeepEPConfig.get_instance().normal_combine_config + or Buffer.get_combine_config(group.size()), + ): + num_nvl_bytes = max( + config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), + num_nvl_bytes, + ) + num_rdma_bytes = max( + config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), + num_rdma_bytes, + ) + if deepep_mode.enable_low_latency(): + assert num_max_dispatch_tokens_per_rank != -1 + assert num_experts != -1 and num_experts % group.size() == 0 + num_rdma_bytes = max( + Buffer.get_low_latency_rdma_size_hint( + num_max_dispatch_tokens_per_rank, + hidden_size, + group.size(), + num_experts, + ), + num_rdma_bytes, + ) + + if deepep_mode == DeepEPMode.NORMAL: + num_qps_per_rank = DeepEPConfig.get_instance().num_sms // 2 + elif deepep_mode in [DeepEPMode.LOW_LATENCY, DeepEPMode.AUTO]: + num_qps_per_rank = num_experts // group.size() + else: + raise NotImplementedError + + if not _is_npu: + total_num_sms = torch.cuda.get_device_properties( + device="cuda" + ).multi_processor_count + if ( + (deepep_mode != DeepEPMode.LOW_LATENCY) + and not is_tbo_enabled() + and (DeepEPConfig.get_instance().num_sms < total_num_sms // 2) + ): + logger.warning( + f"Only use {DeepEPConfig.get_instance().num_sms} SMs for DeepEP communication. " + f"This may result in highly suboptimal performance. " + f"Consider using --deepep-config to change the behavior." + ) + + cls._buffer = Buffer( + group, + num_nvl_bytes, + num_rdma_bytes, + low_latency_mode=deepep_mode.enable_low_latency(), + num_qps_per_rank=num_qps_per_rank, + # TODO can be false when unneeded + allow_mnnvl=True, + ) + return cls._buffer + + @classmethod + def clean_buffer(cls): + if not cls._buffer.low_latency_mode: + return + cls._buffer.clean_low_latency_buffer( + cls._num_max_dispatch_tokens_per_rank, + cls._hidden_size, + cls._num_experts, + ) + + @classmethod + def set_dispatch_mode_as_normal(cls): + cls._dispatch_mode = DeepEPDispatchMode.NORMAL + + @classmethod + def set_dispatch_mode_as_low_latency(cls): + if cls._dispatch_mode == DeepEPDispatchMode.NORMAL: + cls.clean_buffer() + cls._dispatch_mode = DeepEPDispatchMode.LOW_LATENCY + + +class DeepEPConfig(BaseDispatcherConfig): + _instance = None + + def __init__(self): + config_str = get_deepep_config() + if config_str: + config_parsed = load_json_config(config_str) + if torch.distributed.get_rank() == 0: + logger.info(f"Use DeepEP Config: {config_parsed}") + config_dispatch = config_parsed["normal_dispatch"] + config_combine = config_parsed["normal_combine"] + + self.normal_dispatch_config = Config(**config_dispatch) + self.normal_combine_config = Config(**config_combine) + + assert config_dispatch["num_sms"] == config_combine["num_sms"] + self.num_sms = config_dispatch["num_sms"] + else: + self.normal_dispatch_config = None + self.normal_combine_config = None + self.num_sms = Buffer.num_sms + + @classmethod + def get_instance(cls): + if cls._instance is None: + cls._instance = DeepEPConfig() + return cls._instance + + +class _DeepEPDispatcherImplBase: + def __init__( + self, + group: torch.distributed.ProcessGroup, + router_topk: int, + permute_fusion: bool, + num_experts: int, + num_local_experts: int, + hidden_size: int, + params_dtype: torch.dtype, + deepep_mode: DeepEPMode, + ): + if not use_deepep: + raise ImportError( + "DeepEP is not installed. Please install DeepEP package from " + "https://github.com/deepseek-ai/deepep." + ) + + self.group = group + self.router_topk = router_topk + self.permute_fusion = permute_fusion + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.hidden_size = hidden_size + self.params_dtype = params_dtype + self.deepep_mode = deepep_mode + + self.params_bytes = 2 + self.num_max_dispatch_tokens_per_rank = get_int_env_var( + "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 128 + ) + # DeepEP internode_ll dispatch uses FINISHED_SUM_TAG=1024 + # and the logic requires num-tokens-sent-from-one-rank-to-another-rank less than it + assert self.num_max_dispatch_tokens_per_rank <= 1024 + + self.handle = None + + def dispatch_a( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ): + raise NotImplementedError + + def dispatch_b(self, *args, **kwargs): + raise NotImplementedError + + def combine_a( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ): + raise NotImplementedError + + def combine_b(self, *args, **kwargs): + raise NotImplementedError + + def _get_buffer(self): + raise NotImplementedError + + +class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase): + def __init__(self, async_finish: bool, **kwargs): + super().__init__(**kwargs) + + self.async_finish = async_finish + self.src2dst = None + + def dispatch_a( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ): + topk_idx = topk_idx.to(torch.int64) + if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + # TODO hard code 128 block quant,use fp8 communication + hidden_states = sglang_per_token_group_quant_fp8( + hidden_states, + 128, + column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + previous_event = Buffer.capture() if self.async_finish else None + return hidden_states, topk_idx, topk_weights, previous_event + + def dispatch_b(self, hidden_states, topk_idx, topk_weights, previous_event): + ( + hidden_states, + topk_idx, + topk_weights, + num_recv_tokens_per_expert, + event, + ) = self._dispatch_core(hidden_states, topk_idx, topk_weights, previous_event) + event.current_stream_wait() if self.async_finish else () + return DeepEPNormalOutput( + hidden_states, topk_idx, topk_weights, num_recv_tokens_per_expert + ) + + def _dispatch_core( + self, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + previous_event, + ): + buffer = self._get_buffer() + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + previous_event, + ) = buffer.get_dispatch_layout( + topk_idx, + self.num_experts, + previous_event=previous_event, + async_finish=self.async_finish, + allocate_on_comm_stream=previous_event is not None, + ) + + # FIXME: `handle` should be transmitted with tokens from dispatch to combine. + # However, doing this would incur an unknown synchronization error, but keeping + # `handle` as a member variable works. + + ( + recv_x, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert, + self.handle, + event, + ) = buffer.dispatch( + x, + topk_idx=topk_idx, + topk_weights=topk_weights, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=previous_event, + async_finish=self.async_finish, + allocate_on_comm_stream=(previous_event is not None) and self.async_finish, + expert_alignment=128 if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM else 1, + config=DeepEPConfig.get_instance().normal_dispatch_config, + ) + + get_global_expert_distribution_recorder().on_deepep_dispatch_normal( + num_recv_tokens_per_expert, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + num_tokens_per_expert=num_tokens_per_expert, + ) + + return ( + recv_x, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert, + event, + ) + + def combine_a( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ): + from sglang.srt.layers.moe.ep_moe.kernels import ( + deepep_post_reorder_triton_kernel, + ) + + if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM or _use_aiter or _is_npu: + output = hidden_states + else: + if hidden_states.shape[0] > 0: + num_tokens = self.src2dst.shape[0] // self.router_topk + output = torch.empty( + (num_tokens, hidden_states.shape[1]), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + deepep_post_reorder_triton_kernel[(num_tokens,)]( + hidden_states, + output, + self.src2dst, + topk_idx, + topk_weights, + self.router_topk, + hidden_states.shape[1], + BLOCK_SIZE=512, + ) + else: + output = torch.zeros( + (0, hidden_states.shape[1]), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + previous_event = Buffer.capture() if self.async_finish else None + return output, previous_event + + def combine_b(self, output, previous_event): + hidden_states, event = self._combine_core(output, previous_event) + event.current_stream_wait() if self.async_finish else () + self.handle = None + self.src2dst = None + return hidden_states + + def _combine_core(self, x: torch.Tensor, previous_event): + buffer = self._get_buffer() + combined_x, _, event = buffer.combine( + x, + self.handle, + async_finish=self.async_finish, + previous_event=previous_event, + allocate_on_comm_stream=previous_event is not None, + config=DeepEPConfig.get_instance().normal_combine_config, + ) + return combined_x, event + + def _get_buffer(self): + DeepEPBuffer.set_dispatch_mode_as_normal() + + return DeepEPBuffer.get_deepep_buffer( + self.group, + self.hidden_size, + self.params_bytes, + self.deepep_mode, + self.num_max_dispatch_tokens_per_rank, + self.num_experts, + ) + + +class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): + def __init__(self, return_recv_hook: bool, **kwargs): + super().__init__(**kwargs) + + """ + num_max_dispatch_tokens_per_rank: the actual batch size in the decoding engine should be less than 256 + https://github.com/deepseek-ai/DeepEP?tab=readme-ov-file#example-use-in-inference-decoding + """ + self.return_recv_hook = return_recv_hook + + def dispatch_a( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ): + buffer = self._get_buffer() + topk_idx = topk_idx.to(torch.int64) + expected_m = ( + hidden_states.shape[0] * buffer.group_size * topk_idx.shape[1] + + self.num_experts + ) // self.num_experts + hidden_states, masked_m, event, hook = self._dispatch_core( + hidden_states, + topk_idx, + use_fp8=True, + ) + return ( + hidden_states, + topk_idx, + topk_weights, + masked_m, + expected_m, + event, + hook, + ) + + def dispatch_b( + self, + hidden_states, + topk_idx, + topk_weights, + masked_m, + expected_m, + event, + hook, + ): + hook() if self.return_recv_hook else event.current_stream_wait() + + get_global_expert_distribution_recorder().on_deepep_dispatch_low_latency( + masked_m + ) + + deepep_output = DeepEPLLOutput( + hidden_states, + topk_idx, + topk_weights, + masked_m, + expected_m, + ) + return deepep_output + + def _dispatch_core( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + use_fp8: bool = False, + ): + buffer = self._get_buffer() + packed_recv_hidden, packed_recv_count, self.handle, event, hook = ( + buffer.low_latency_dispatch( + hidden_states, + topk_idx, + self.num_max_dispatch_tokens_per_rank, + self.num_experts, + use_fp8=use_fp8, + async_finish=not self.return_recv_hook, + return_recv_hook=self.return_recv_hook, + round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and deep_gemm_wrapper.DEEPGEMM_BLACKWELL, + use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and deep_gemm_wrapper.DEEPGEMM_BLACKWELL, + ) + ) + return packed_recv_hidden, packed_recv_count, event, hook + + def combine_a( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ): + hidden_states, event, hook = self._combine_core( + hidden_states, + topk_idx, + topk_weights, + ) + return hidden_states, event, hook + + def combine_b(self, hidden_states, event, hook): + hook() if self.return_recv_hook else event.current_stream_wait() + return hidden_states + + def _combine_core( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ): + buffer = self._get_buffer() + combined_hidden_states, event, hook = buffer.low_latency_combine( + hidden_states, + topk_idx, + topk_weights, + self.handle, + async_finish=not self.return_recv_hook, + return_recv_hook=self.return_recv_hook, + ) + self.handle = None + return combined_hidden_states, event, hook + + def _get_buffer(self): + DeepEPBuffer.set_dispatch_mode_as_low_latency() + return DeepEPBuffer.get_deepep_buffer( + self.group, + self.hidden_size, + self.params_bytes, + self.deepep_mode, + self.num_max_dispatch_tokens_per_rank, + self.num_experts, + ) + + +@dataclass +class _Stage(Enum): + INITIAL = auto() + AFTER_DISPATCH_A = auto() + AFTER_DISPATCH_B = auto() + AFTER_COMBINE_A = auto() + + +class DeepEPDispatcher(BaseDispatcher): + def __init__( + self, + group: torch.distributed.ProcessGroup, + router_topk: int, + permute_fusion: bool = False, + num_experts: int = None, + num_local_experts: int = None, + hidden_size: int = None, + params_dtype: torch.dtype = None, + deepep_mode: DeepEPMode = DeepEPMode.AUTO, + async_finish: bool = False, + return_recv_hook: bool = False, + ): + self.deepep_mode = deepep_mode + + common_kwargs = dict( + group=group, + router_topk=router_topk, + permute_fusion=permute_fusion, + num_experts=num_experts, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + params_dtype=params_dtype, + deepep_mode=deepep_mode, + ) + + if self.deepep_mode.enable_low_latency(): + self._low_latency_dispatcher = _DeepEPDispatcherImplLowLatency( + return_recv_hook=return_recv_hook, + **common_kwargs, + ) + if self.deepep_mode.enable_normal(): + self._normal_dispatcher = _DeepEPDispatcherImplNormal( + async_finish=async_finish, + **common_kwargs, + ) + + self._stage = _Stage.INITIAL + + def dispatch(self, *args, **kwargs) -> DispatchOutput: + self.dispatch_a(*args, **kwargs) + ret = self.dispatch_b() + return ret + + def dispatch_a( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + forward_batch: ForwardBatch, + ): + self._update_stage(_Stage.INITIAL, _Stage.AFTER_DISPATCH_A) + inner_state = self._get_impl(forward_batch).dispatch_a( + hidden_states=hidden_states, + topk_idx=topk_idx, + topk_weights=topk_weights, + ) + self._dispatch_intermediate_state = forward_batch, inner_state + + def dispatch_b(self): + self._update_stage(_Stage.AFTER_DISPATCH_A, _Stage.AFTER_DISPATCH_B) + forward_batch, inner_state = self._dispatch_intermediate_state + del self._dispatch_intermediate_state + return self._get_impl(forward_batch).dispatch_b(*inner_state) + + def combine(self, *args, **kwargs) -> Tuple: + self.combine_a(*args, **kwargs) + ret = self.combine_b() + return ret + + def combine_a( + self, + hidden_states: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + forward_batch: ForwardBatch, + ): + self._update_stage(_Stage.AFTER_DISPATCH_B, _Stage.AFTER_COMBINE_A) + inner_state = self._get_impl(forward_batch).combine_a( + hidden_states=hidden_states, + topk_idx=topk_idx, + topk_weights=topk_weights, + ) + self._combine_intermediate_state = forward_batch, inner_state + + def combine_b(self): + self._update_stage(_Stage.AFTER_COMBINE_A, _Stage.INITIAL) + forward_batch, inner_state = self._combine_intermediate_state + del self._combine_intermediate_state + return self._get_impl(forward_batch).combine_b(*inner_state) + + def _get_impl(self, forward_batch: ForwardBatch) -> _DeepEPDispatcherImplBase: + resolved_deepep_mode = self.deepep_mode.resolve( + forward_batch.is_extend_in_batch + ) + if resolved_deepep_mode == DeepEPMode.NORMAL: + return self._normal_dispatcher + elif resolved_deepep_mode == DeepEPMode.LOW_LATENCY: + return self._low_latency_dispatcher + else: + raise ValueError(f"Invalid deepep_mode: {self.deepep_mode}") + + def _update_stage(self, old_stage, new_stage): + assert self._stage == old_stage + self._stage = new_stage diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py new file mode 100644 index 000000000..f984104f6 --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/standard.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple + +import torch + +from sglang.srt.layers.moe.token_dispatcher.base import ( + BaseDispatcher, + CombineInput, + CombineInputFormat, + DispatchOutput, + DispatchOutputFormat, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.topk import TopKOutput + + +class StandardDispatchOutput(NamedTuple): + """Standard dispatch output.""" + + hidden_states: torch.Tensor + topk_output: TopKOutput + + @property + def format(self) -> DispatchOutputFormat: + return DispatchOutputFormat.STANDARD + + +assert isinstance(StandardDispatchOutput, DispatchOutput) + + +class StandardCombineInput(NamedTuple): + """Standard combine input.""" + + hidden_states: torch.Tensor + + @property + def format(self) -> CombineInputFormat: + return CombineInputFormat.STANDARD + + +assert isinstance(StandardCombineInput, CombineInput) + + +class StandardDispatcher(BaseDispatcher): + + def dispatch( + self, hidden_states: torch.Tensor, topk_output: TopKOutput + ) -> DispatchOutput: + return StandardDispatchOutput( + hidden_states=hidden_states, topk_output=topk_output + ) + + def combine(self, combine_input: CombineInput) -> torch.Tensor: + if isinstance(combine_input, StandardCombineInput): + return combine_input.hidden_states + else: + # TODO: this branch should be removed in the future + assert isinstance(combine_input, torch.Tensor) + return combine_input diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py new file mode 100644 index 000000000..b8f73473c --- /dev/null +++ b/python/sglang/srt/layers/moe/topk.py @@ -0,0 +1,876 @@ +# Copyright 2024 SGLang 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 __future__ import annotations + +import logging +import math +from dataclasses import dataclass +from enum import Enum, auto +from typing import ( + Callable, + NamedTuple, + Optional, + Protocol, + TypeGuard, + runtime_checkable, +) + +import torch +import torch.nn.functional as F + +from sglang.srt.custom_op import CustomOp +from sglang.srt.eplb import expert_location_dispatch +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location_dispatch import ( + ExpertLocationDispatchInfo, + topk_ids_logical_to_physical, +) +from sglang.srt.layers.moe import ( + get_moe_runner_backend, + should_use_flashinfer_trtllm_moe, +) +from sglang.srt.utils import ( + cpu_has_amx_support, + get_bool_env_var, + get_compiler_backend, + is_cpu, + is_cuda, + is_hip, + is_npu, +) + +try: + from triton_kernels.routing import GatherIndx, RoutingData, ScatterIndx, routing +except ImportError: + pass +logger = logging.getLogger(__name__) + + +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_cpu = is_cpu() +_is_cpu_amx_available = cpu_has_amx_support() +_is_npu = is_npu() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _is_cuda: + from sgl_kernel import moe_fused_gate + +if _is_cuda or _is_hip: + from sgl_kernel import topk_softmax +if _use_aiter: + try: + from aiter import biased_grouped_topk as aiter_biased_grouped_topk + except ImportError: + raise ImportError("aiter is required when SGLANG_USE_AITER is set to True") +if _is_npu: + import torch_npu + +# -------------------------------- TopKConfig --------------------------------------- + + +@dataclass +class TopKConfig: + top_k: int + use_grouped_topk: bool = False + topk_group: Optional[int] = None + num_expert_group: Optional[int] = None + renormalize: bool = True + num_fused_shared_experts: int = 0 + custom_routing_function: Optional[Callable] = None + correction_bias: Optional[torch.Tensor] = None + torch_native: bool = False + routed_scaling_factor: Optional[float] = None + apply_routed_scaling_factor_on_output: bool = False + + +# -------------------------------- TopKOutput --------------------------------------- + + +class TopKOutputChecker: + + @staticmethod + def format_is_standard(topk_output: TopKOutput) -> TypeGuard[StandardTopKOutput]: + return topk_output.format.is_standard() + + @staticmethod + def format_is_triton_kernel( + topk_output: TopKOutput, + ) -> TypeGuard[TritonKernelTopKOutput]: + return topk_output.format.is_triton_kernel() + + @staticmethod + def format_is_bypassed(topk_output: TopKOutput) -> TypeGuard[BypassedTopKOutput]: + return topk_output.format.is_bypassed() + + +class TopKOutputFormat(Enum): + STANDARD = auto() + TRITON_KERNEL = auto() + BYPASSED = auto() + + def is_standard(self) -> bool: + return self == TopKOutputFormat.STANDARD + + def is_triton_kernel(self) -> bool: + return self == TopKOutputFormat.TRITON_KERNEL + + def is_bypassed(self) -> bool: + return self == TopKOutputFormat.BYPASSED + + +@runtime_checkable +class TopKOutput(Protocol): + """Protocol for top-k outputs in different formats.""" + + @property + def format(self) -> TopKOutputFormat: + """The format of the output.""" + ... + + +class StandardTopKOutput(NamedTuple): + """Standard top-k output format.""" + + topk_weights: torch.Tensor + topk_ids: torch.Tensor + router_logits: torch.Tensor + + @property + def format(self) -> TopKOutputFormat: + return TopKOutputFormat.STANDARD + + +class TritonKernelTopKOutput(NamedTuple): + """Triton kernel top-k output format.""" + + routing_data: RoutingData + gather_indx: GatherIndx + scatter_indx: ScatterIndx + + @property + def format(self) -> TopKOutputFormat: + return TopKOutputFormat.TRITON_KERNEL + + +class BypassedTopKOutput(NamedTuple): + """Bypassed top-k output format.""" + + hidden_states: torch.Tensor + router_logits: torch.Tensor + topk_config: TopKConfig + num_token_non_padded: Optional[torch.Tensor] = None + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None + + @property + def format(self) -> TopKOutputFormat: + return TopKOutputFormat.BYPASSED + + +# -------------------------------- TopK --------------------------------------- + + +class TopK(CustomOp): + + def __init__( + self, + top_k: int, + *, + use_grouped_topk: bool = False, + topk_group: Optional[int] = None, + num_expert_group: Optional[int] = None, + renormalize: bool = True, + num_fused_shared_experts: int = 0, + custom_routing_function: Optional[Callable] = None, + scoring_func: str = "softmax", + correction_bias: Optional[torch.Tensor] = None, + routed_scaling_factor: Optional[float] = None, + apply_routed_scaling_factor_on_output: Optional[bool] = False, + force_topk: bool = False, + ): + # NOTE: scoring_func is not used for now, but we keep it for future use + # see https://github.com/sgl-project/sglang/pull/4505 for more details + super().__init__() + + if use_grouped_topk: + assert num_expert_group is not None and topk_group is not None + + self.topk_config = TopKConfig( + top_k=top_k, + use_grouped_topk=use_grouped_topk, + renormalize=renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + num_fused_shared_experts=num_fused_shared_experts, + custom_routing_function=custom_routing_function, + correction_bias=correction_bias, + routed_scaling_factor=routed_scaling_factor, + apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, + ) + + self.use_triton_kernels = get_moe_runner_backend().is_triton_kernel() + self.force_topk = force_topk + + def forward_native( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + *, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + ) -> TopKOutput: + self.topk_config.torch_native = True + return select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=self.topk_config, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) + + def forward_cuda( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + *, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + ) -> TopKOutput: + if self.use_triton_kernels: + # renormalize=True is equivalent to sm_first=False + routing_data, gather_idx, scatter_idx = routing( + router_logits, + self.topk_config.top_k, + sm_first=not self.topk_config.renormalize, + ) + return TritonKernelTopKOutput(routing_data, gather_idx, scatter_idx) + elif not self.force_topk and ( + should_use_flashinfer_trtllm_moe() + or get_moe_runner_backend().is_flashinfer_mxfp4() + ): + return BypassedTopKOutput( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=self.topk_config, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) + else: + self.topk_config.torch_native = False + return select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=self.topk_config, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) + + def forward_cpu( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + *, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + ) -> TopKOutput: + return select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=self.topk_config, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) + + def forward_npu( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + *, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + ) -> TopKOutput: + global_num_experts = router_logits.shape[-1] + + # NOTE: now npu_moe_gating_top_k can only support `group_count=256` pattern + if global_num_experts == 256: + + routed_scaling_factor = self.topk_config.routed_scaling_factor or 1 + router_logits = router_logits.to(torch.float32) + + topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k( + router_logits, + k=self.topk_config.top_k, + bias=self.topk_config.correction_bias.to(torch.float32), + k_group=self.topk_config.topk_group, + group_count=self.topk_config.num_expert_group, + group_select_mode=1, + renorm=0, + norm_type=1, + routed_scaling_factor=routed_scaling_factor, + eps=float(1e-20), + ) + + if self.topk_config.renormalize: + topk_weights_sum = ( + topk_weights.sum(dim=-1, keepdim=True) + if self.topk_config.num_fused_shared_experts == 0 + else topk_weights[:, :-1].sum(dim=-1, keepdim=True) + ) + topk_weights = topk_weights / topk_weights_sum + + if expert_location_dispatch_info is not None: + topk_ids = topk_ids_logical_to_physical( + topk_ids, expert_location_dispatch_info + ) + get_global_expert_distribution_recorder().on_select_experts( + topk_ids=topk_ids + ) + + return StandardTopKOutput(topk_weights, topk_ids, _) + else: + self.topk_config.torch_native = True + return select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=self.topk_config, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) + + def empty_topk_output(self, device: torch.device) -> TopKOutput: + topk = self.topk_config.top_k - self.topk_config.num_fused_shared_experts + topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device) + topk_idx = torch.full((0, topk), -1, dtype=torch.int32, device=device) + router_logits = torch.empty((0, topk), dtype=torch.float32, device=device) + return StandardTopKOutput(topk_weights, topk_idx, router_logits) + + +# ------------------------------- TopK implementation ------------------------------------- + + +def fused_topk_torch_native( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + correction_bias: torch.Tensor = None, +): + if correction_bias is not None: + n_routed_experts = gating_output.shape[-1] + scores = gating_output.softmax(dim=-1) + scores_for_choice = scores.view( + -1, n_routed_experts + ) + correction_bias.unsqueeze(0) + topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=False)[1] + topk_weights = scores.gather(1, topk_ids) + else: + assert ( + hidden_states.shape[0] == gating_output.shape[0] + ), f"Number of tokens mismatch, {hidden_states.shape=} vs {gating_output.shape=}" + M, _ = hidden_states.shape + topk_weights = torch.empty( + M, topk, dtype=torch.float32, device=hidden_states.device + ) + topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device) + topk_weights = F.softmax(gating_output.float(), dim=-1) + topk_weights, topk_ids = torch.topk(topk_weights, topk, dim=-1) + + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + return topk_weights, topk_ids + + +def fused_topk_cpu( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + correction_bias: torch.Tensor = None, +): + topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu( + hidden_states=hidden_states, + gating_output=gating_output, + topk=topk, + renormalize=renormalize, + ) + topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) + _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + return topk_weights, topk_ids + + +def apply_topk_weights_cpu(need_apply, topk_weights, inputs): + if not need_apply: + return inputs, topk_weights + + # TODO: fuse below processing in fused_experts_cpu kernel + inputs = inputs * topk_weights.to(inputs.dtype) + topk_weights = torch.ones_like( + topk_weights, dtype=torch.float32 + ) # clear topk_weights as already applied + + return inputs, topk_weights + + +def fused_topk( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, +): + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + + M, _ = hidden_states.shape + + topk_weights = torch.empty( + M, topk, dtype=torch.float32, device=hidden_states.device + ) + topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device) + + topk_softmax( + topk_weights, + topk_ids, + gating_output, + renormalize, + ) + + topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) + _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + return topk_weights, topk_ids + + +# This is used by the Deepseek V2/V3/R1 series models +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def grouped_topk_gpu( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: Optional[int] = None, + topk_group: Optional[int] = None, + num_fused_shared_experts: int = 0, + routed_scaling_factor: Optional[float] = None, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + apply_routed_scaling_factor_on_output: Optional[bool] = False, +): + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + + scores = torch.softmax(gating_output, dim=-1) + # NPU compiler limitation + if _is_npu and scores.dtype == torch.bfloat16: + scores = scores.to(torch.float16) + num_token = scores.shape[0] + num_experts = scores.shape[1] + group_scores = ( + scores.view(num_token, num_expert_group, -1).max(dim=-1).values + ) # [n, n_group] + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[ + 1 + ] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) + .reshape(num_token, -1) + ) # [n, e] + tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e] + # TODO: NPU can't support directly evaluating a comparison for now + topk_weights, topk_ids = torch.topk( + tmp_scores, + k=topk, + dim=-1, + sorted=(True if num_fused_shared_experts > 0 else False), + ) + if num_fused_shared_experts: + topk_ids[:, -1] = torch.randint( + low=num_experts, + high=num_experts + num_fused_shared_experts, + size=(topk_ids.size(0),), + dtype=topk_ids.dtype, + device=topk_ids.device, + ) + topk_weights[:, -1] = topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor + + if renormalize: + topk_weights_sum = ( + topk_weights.sum(dim=-1, keepdim=True) + if num_fused_shared_experts == 0 + else topk_weights[:, :-1].sum(dim=-1, keepdim=True) + ) + topk_weights = topk_weights / topk_weights_sum + if apply_routed_scaling_factor_on_output: + topk_weights *= routed_scaling_factor + + topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32) + topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) + _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + return topk_weights, topk_ids + + +def grouped_topk_cpu( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: Optional[int] = None, + topk_group: Optional[int] = None, + num_fused_shared_experts: int = 0, + routed_scaling_factor: Optional[float] = None, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + apply_routed_scaling_factor_on_output: Optional[bool] = False, +): + assert not apply_routed_scaling_factor_on_output + assert expert_location_dispatch_info is None + return torch.ops.sgl_kernel.grouped_topk_cpu( + hidden_states, + gating_output, + topk, + renormalize, + num_expert_group, + topk_group, + num_fused_shared_experts, + routed_scaling_factor, + num_token_non_padded, + ) + + +@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu) +def biased_grouped_topk_impl( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + correction_bias: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: Optional[int] = None, + topk_group: Optional[int] = None, + num_fused_shared_experts: int = 0, + routed_scaling_factor: Optional[float] = None, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + apply_routed_scaling_factor_on_output: Optional[bool] = False, +): + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + + scores = gating_output.sigmoid() + num_token = scores.shape[0] + num_experts = scores.shape[1] + scores_for_choice = scores.view(num_token, -1) + correction_bias.unsqueeze(0) + group_scores = ( + scores_for_choice.view(num_token, num_expert_group, -1) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) # [n, n_group] + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[ + 1 + ] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) + .reshape(num_token, -1) + ) # [n, e] + tmp_scores = scores_for_choice.masked_fill( + ~score_mask.bool(), float("-inf") + ) # [n, e] + # TODO: NPU can't support directly evaluating a comparison for now + _, topk_ids = torch.topk( + tmp_scores, + k=topk, + dim=-1, + sorted=(True if num_fused_shared_experts > 0 else False), + ) + topk_weights = scores.gather(1, topk_ids) + + if num_fused_shared_experts: + topk_ids[:, -1] = torch.randint( + low=num_experts, + high=num_experts + num_fused_shared_experts, + size=(topk_ids.size(0),), + dtype=topk_ids.dtype, + device=topk_ids.device, + ) + topk_weights[:, -1] = topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor + + if renormalize: + topk_weights_sum = ( + topk_weights.sum(dim=-1, keepdim=True) + if num_fused_shared_experts == 0 + else topk_weights[:, :-1].sum(dim=-1, keepdim=True) + ) + topk_weights = topk_weights / topk_weights_sum + if apply_routed_scaling_factor_on_output: + topk_weights *= routed_scaling_factor + + topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32) + topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) + _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + return topk_weights, topk_ids + + +def is_power_of_two(n): + return n > 0 and math.log2(n).is_integer() + + +def _mask_topk_ids_padded_region( + topk_ids: torch.Tensor, + num_token_non_padded: Optional[torch.Tensor] = None, +): + if num_token_non_padded is None: + return + indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device) + topk_ids[indices >= num_token_non_padded, :] = -1 + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def _biased_grouped_topk_postprocess( + topk_ids, expert_location_dispatch_info, num_token_non_padded +): + topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) + _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + return topk_ids + + +def biased_grouped_topk_gpu( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + correction_bias: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: Optional[int] = None, + topk_group: Optional[int] = None, + num_fused_shared_experts: int = 0, + routed_scaling_factor: Optional[float] = None, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + apply_routed_scaling_factor_on_output: Optional[bool] = False, +): + assert ( + routed_scaling_factor is not None + ), "routed_scaling_factor is required for biased_grouped_topk" + # TODO: moe_fused_gate kernel is not supported for num_fused_shared_experts > 0 now. + if ( + _is_cuda + and gating_output.shape[1] // num_expert_group + <= 32 # moe_fused_gate kernel ensure that num_experts/num_expert_group does not exceed MAX_VPT=32 now. And when kernel can handle MAX_VPT > 32, we can remove this assertion. + and is_power_of_two(correction_bias.shape[0]) + ): + topk_weights, topk_ids = moe_fused_gate( + gating_output.to(dtype=torch.float32), + correction_bias, + num_expert_group, + topk_group, + topk, + num_fused_shared_experts, + routed_scaling_factor, + apply_routed_scaling_factor_on_output, + ) + # TODO merge into kernel + if (expert_location_dispatch_info is not None) or ( + num_token_non_padded is not None + ): + topk_ids = _biased_grouped_topk_postprocess( + topk_ids, expert_location_dispatch_info, num_token_non_padded + ) + return topk_weights, topk_ids + elif _use_aiter: + assert not apply_routed_scaling_factor_on_output, "Not implemented" + token = gating_output.shape[0] + device = gating_output.device + assert ( + hidden_states.shape[0] == gating_output.shape[0] + ), f"Number of tokens mismatch: hidden_states.shape[0] = {hidden_states.shape[0]}, gating_output.shape[0] = {gating_output.shape[0]}" + topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device) + topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) + aiter_biased_grouped_topk( + gating_output.to(dtype=torch.float32), + correction_bias, + topk_weights, + topk_ids, + num_expert_group, + topk_group, + renormalize, + routed_scaling_factor, + ) + return topk_weights, topk_ids + else: + return biased_grouped_topk_impl( + hidden_states, + gating_output, + correction_bias, + topk, + renormalize, + num_expert_group, + topk_group, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, + ) + + +def biased_grouped_topk_cpu( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + correction_bias: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: Optional[int] = None, + topk_group: Optional[int] = None, + compiled: bool = True, + num_fused_shared_experts: int = 0, + routed_scaling_factor: Optional[float] = None, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + apply_routed_scaling_factor_on_output: Optional[bool] = False, +): + assert expert_location_dispatch_info is None + assert not apply_routed_scaling_factor_on_output, "Not implemented" + return torch.ops.sgl_kernel.biased_grouped_topk_cpu( + hidden_states, + gating_output, + correction_bias, + topk, + renormalize, + num_expert_group, + topk_group, + num_fused_shared_experts, + routed_scaling_factor, + num_token_non_padded, + ) + + +if _is_cpu and _is_cpu_amx_available: + biased_grouped_topk = biased_grouped_topk_cpu + grouped_topk = grouped_topk_cpu + fused_topk_native = fused_topk_cpu + fused_topk = fused_topk_cpu +else: + biased_grouped_topk = biased_grouped_topk_gpu + grouped_topk = grouped_topk_gpu + fused_topk_native = fused_topk_torch_native + + +def select_experts( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + topk_config: TopKConfig, + *, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, +) -> StandardTopKOutput: + + top_k = topk_config.top_k + use_grouped_topk = topk_config.use_grouped_topk + topk_group = topk_config.topk_group + num_expert_group = topk_config.num_expert_group + renormalize = topk_config.renormalize + num_fused_shared_experts = topk_config.num_fused_shared_experts + custom_routing_function = topk_config.custom_routing_function + correction_bias = topk_config.correction_bias + torch_native = topk_config.torch_native + routed_scaling_factor = topk_config.routed_scaling_factor + apply_routed_scaling_factor_on_output = ( + topk_config.apply_routed_scaling_factor_on_output + ) + + router_logits, correction_bias = ( + expert_location_dispatch.transform_select_experts_inputs( + router_logits=router_logits, + correction_bias=correction_bias, + info=expert_location_dispatch_info, + ) + ) + + # DeepSeek V2/V3/R1 series models use grouped_top_k + if use_grouped_topk: + assert topk_group is not None + assert num_expert_group is not None + if correction_bias is None: + topk_weights, topk_ids = grouped_topk( + hidden_states=hidden_states, + gating_output=router_logits, + topk=top_k, + renormalize=renormalize, + num_expert_group=num_expert_group, + topk_group=topk_group, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, + ) + else: + topk_weights, topk_ids = biased_grouped_topk( + hidden_states=hidden_states, + gating_output=router_logits, + correction_bias=correction_bias, + topk=top_k, + renormalize=renormalize, + num_expert_group=num_expert_group, + topk_group=topk_group, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, + ) + elif torch_native and custom_routing_function is None: + assert ( + num_token_non_padded is None + ), "num_token_non_padded is not yet supported in fused_topk_native" + assert expert_location_dispatch_info is None + assert not apply_routed_scaling_factor_on_output, "Not implemented" + topk_weights, topk_ids = fused_topk_native( + hidden_states=hidden_states, + gating_output=router_logits, + topk=top_k, + renormalize=renormalize, + correction_bias=correction_bias, + ) + elif custom_routing_function is None: + assert not apply_routed_scaling_factor_on_output, "Not implemented" + # Qwen3MOE uses fused_topk + topk_weights, topk_ids = fused_topk( + hidden_states=hidden_states, + gating_output=router_logits, + topk=top_k, + renormalize=renormalize, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) + else: + assert ( + num_token_non_padded is None + ), "num_token_non_padded is not yet supported in custom_routing_function" + assert expert_location_dispatch_info is None + assert not apply_routed_scaling_factor_on_output, "Not implemented" + topk_weights, topk_ids = custom_routing_function( + hidden_states=hidden_states, + gating_output=router_logits, + topk=top_k, + renormalize=renormalize, + ) + + get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids) + + return StandardTopKOutput(topk_weights, topk_ids, router_logits) diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py new file mode 100644 index 000000000..b4e4ec424 --- /dev/null +++ b/python/sglang/srt/layers/moe/utils.py @@ -0,0 +1,201 @@ +from __future__ import annotations + +import importlib.util +import logging +from enum import Enum +from functools import lru_cache +from typing import TYPE_CHECKING, Optional + +from packaging import version as pkg_version + +from sglang.srt.distributed.parallel_state import get_moe_expert_parallel_world_size +from sglang.srt.layers.dp_attention import ( + get_attention_dp_size, + is_dp_attention_enabled, +) + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +class MoeA2ABackend(Enum): + + NONE = "none" + DEEPEP = "deepep" + + @classmethod + def _missing_(cls, value): + if value is None: + return cls.NONE + for member in cls: + if value == member.value: + return member + raise ValueError(f"No {cls.__name__} member for value {value}") + + def is_none(self): + return self == MoeA2ABackend.NONE + + def is_deepep(self): + return self == MoeA2ABackend.DEEPEP + + +class MoeRunnerBackend(Enum): + + AUTO = "auto" + TRITON = "triton" + TRITON_KERNEL = "triton_kernel" + FLASHINFER = "flashinfer_trtllm" + FLASHINFER_CUTLASS = "flashinfer_cutlass" + FLASHINFER_MXFP4 = "flashinfer_mxfp4" + + def is_auto(self): + return self == MoeRunnerBackend.AUTO + + def is_triton(self): + return self == MoeRunnerBackend.TRITON + + def is_triton_kernel(self): + return self == MoeRunnerBackend.TRITON_KERNEL + + def is_flashinfer_trtllm(self): + return self == MoeRunnerBackend.FLASHINFER + + def is_flashinfer_cutlass(self): + return self == MoeRunnerBackend.FLASHINFER_CUTLASS + + def is_flashinfer_mxfp4(self): + return self == MoeRunnerBackend.FLASHINFER_MXFP4 + + +class DeepEPMode(Enum): + + NORMAL = "normal" + LOW_LATENCY = "low_latency" + AUTO = "auto" + + def enable_normal(self) -> bool: + return self in [DeepEPMode.NORMAL, DeepEPMode.AUTO] + + def enable_low_latency(self) -> bool: + return self in [DeepEPMode.LOW_LATENCY, DeepEPMode.AUTO] + + def resolve(self, is_extend_in_batch: bool) -> DeepEPMode: + if self != DeepEPMode.AUTO: + return self + + if is_extend_in_batch: + return DeepEPMode.NORMAL + else: + return DeepEPMode.LOW_LATENCY + + def is_normal(self) -> bool: + return self == DeepEPMode.NORMAL + + def is_low_latency(self) -> bool: + return self == DeepEPMode.LOW_LATENCY + + def is_auto(self) -> bool: + return self == DeepEPMode.AUTO + + +MOE_A2A_BACKEND: Optional[MoeA2ABackend] = None +MOE_RUNNER_BACKEND: Optional[MoeRunnerBackend] = None +DEEPEP_MODE: Optional[DeepEPMode] = None +IS_TBO_ENABLED: Optional[bool] = None +TBO_TOKEN_DISTRIBUTION_THRESHOLD: Optional[float] = None +DEEPEP_CONFIG: Optional[str] = None +DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER: Optional[bool] = None + + +def initialize_moe_config(server_args: ServerArgs): + global MOE_A2A_BACKEND + global MOE_RUNNER_BACKEND + global DEEPEP_MODE + global DEEPEP_CONFIG + global IS_TBO_ENABLED + global TBO_TOKEN_DISTRIBUTION_THRESHOLD + global DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER + + MOE_A2A_BACKEND = MoeA2ABackend(server_args.moe_a2a_backend) + MOE_RUNNER_BACKEND = MoeRunnerBackend(server_args.moe_runner_backend) + DEEPEP_MODE = DeepEPMode(server_args.deepep_mode) + DEEPEP_CONFIG = server_args.deepep_config or "" + IS_TBO_ENABLED = server_args.enable_two_batch_overlap + TBO_TOKEN_DISTRIBUTION_THRESHOLD = server_args.tbo_token_distribution_threshold + DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER = ( + server_args.disable_flashinfer_cutlass_moe_fp4_allgather + ) + + +def get_moe_a2a_backend() -> MoeA2ABackend: + global MOE_A2A_BACKEND + if MOE_A2A_BACKEND is None: + logger.warning("MOE_A2A_BACKEND is not initialized, using default backend") + MOE_A2A_BACKEND = MoeA2ABackend.NONE + return MOE_A2A_BACKEND + + +def get_moe_runner_backend() -> MoeRunnerBackend: + global MOE_RUNNER_BACKEND + if MOE_RUNNER_BACKEND is None: + logger.warning("MOE_RUNNER_BACKEND is not initialized, using triton backend") + MOE_RUNNER_BACKEND = MoeRunnerBackend.AUTO + return MOE_RUNNER_BACKEND + + +def get_deepep_mode() -> DeepEPMode: + global DEEPEP_MODE + if DEEPEP_MODE is None: + logger.warning("DEEPEP_MODE is not initialized, using auto mode") + DEEPEP_MODE = DeepEPMode.AUTO + return DEEPEP_MODE + + +def get_deepep_config() -> str: + global DEEPEP_CONFIG + if DEEPEP_CONFIG is None: + logger.warning("DEEPEP_CONFIG is not initialized, using default config") + DEEPEP_CONFIG = "" + return DEEPEP_CONFIG + + +def is_tbo_enabled() -> bool: + global IS_TBO_ENABLED + if IS_TBO_ENABLED is None: + IS_TBO_ENABLED = False + return IS_TBO_ENABLED + + +def get_tbo_token_distribution_threshold() -> float: + global TBO_TOKEN_DISTRIBUTION_THRESHOLD + if TBO_TOKEN_DISTRIBUTION_THRESHOLD is None: + logger.warning( + "TBO_TOKEN_DISTRIBUTION_THRESHOLD is not initialized, using 0.48" + ) + TBO_TOKEN_DISTRIBUTION_THRESHOLD = 0.48 + return TBO_TOKEN_DISTRIBUTION_THRESHOLD + + +@lru_cache(maxsize=1) +def should_use_flashinfer_trtllm_moe(): + result = get_moe_runner_backend().is_flashinfer_trtllm() and ( + not importlib.util.find_spec("flashinfer") + or pkg_version.parse(__import__("flashinfer").__version__) + >= pkg_version.parse("0.2.9rc1") + ) + return result + + +@lru_cache(maxsize=1) +def should_use_flashinfer_cutlass_moe_fp4_allgather(): + """ + Perform FP4 quantize before all-gather for flashinfer cutlass moe to reduce communication cost for high-throughput serving. + """ + return ( + not DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER + and get_moe_runner_backend().is_flashinfer_cutlass() + and is_dp_attention_enabled() + and get_moe_expert_parallel_world_size() == get_attention_dp_size() + ) diff --git a/python/sglang/srt/layers/multimodal.py b/python/sglang/srt/layers/multimodal.py new file mode 100644 index 000000000..738c65830 --- /dev/null +++ b/python/sglang/srt/layers/multimodal.py @@ -0,0 +1,189 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Logits processing.""" + +import torch +import triton +import triton.language as tl + +FMIX32_C1 = 0x85EBCA6B +FMIX32_C2 = 0xC2B2AE35 +POS_C1 = 0x27D4EB2D +POS_C2 = 0x165667B1 + + +@triton.jit +def _rotl32(x, r: tl.constexpr): + return (x << r) | (x >> (32 - r)) + + +@triton.jit +def _fmix32(x, C1: tl.constexpr, C2: tl.constexpr): + c1 = tl.full((), C1, tl.uint32) + c2 = tl.full((), C2, tl.uint32) + x ^= x >> 16 + x = x * c1 + x ^= x >> 13 + x = x * c2 + x ^= x >> 16 + return x + + +@triton.jit +def hash_tiles32_kernel_blocked( + in_ptr, + out_ptr, + n_u32, + seed1, + seed2, + FM_C1: tl.constexpr, + FM_C2: tl.constexpr, + POS_A: tl.constexpr, + POS_B: tl.constexpr, + TILE: tl.constexpr, + BLOCK: tl.constexpr, + USE_CG: tl.constexpr, +): + pid = tl.program_id(axis=0) + base = pid * TILE + + s1 = tl.full((), seed1, tl.uint32) + s2 = tl.full((), seed2, tl.uint32) + posA = tl.full((), POS_A, tl.uint32) + posB = tl.full((), POS_B, tl.uint32) + + h1 = tl.zeros((), dtype=tl.uint32) + h2 = tl.zeros((), dtype=tl.uint32) + + for off in tl.static_range(0, TILE, BLOCK): + idx = base + off + tl.arange(0, BLOCK) + m = idx < n_u32 + + if USE_CG: + v = tl.load(in_ptr + idx, mask=m, other=0, cache_modifier=".cg") + else: + v = tl.load(in_ptr + idx, mask=m, other=0) + v = v.to(tl.uint32) + + iu = idx.to(tl.uint32) + p1 = (iu * posA + s1) ^ _rotl32(iu, 15) + p2 = (iu * posB + s2) ^ _rotl32(iu, 13) + + k1 = _fmix32(v ^ p1, C1=FM_C1, C2=FM_C2) + k2 = _fmix32(v ^ p2, C1=FM_C1, C2=FM_C2) + + zero32 = tl.zeros_like(k1) + k1 = tl.where(m, k1, zero32) + k2 = tl.where(m, k2, zero32) + + h1 += tl.sum(k1, axis=0).to(tl.uint32) + h2 += tl.sum(k2, axis=0).to(tl.uint32) + + nbytes = tl.full((), n_u32 * 4, tl.uint32) + h1 ^= nbytes + h2 ^= nbytes + h1 = _fmix32(h1, C1=FM_C1, C2=FM_C2) + h2 = ( + _fmix32(h2, C1=FMIX32_C1, C2=FMIX32_C2) + if False + else _fmix32(h2, C1=FM_C1, C2=FM_C2) + ) + + out = (h1.to(tl.uint64) << 32) | h2.to(tl.uint64) + tl.store(out_ptr + pid, out) + + +@triton.jit +def add_tree_reduce_u64_kernel(in_ptr, out_ptr, n_elems, CHUNK: tl.constexpr): + pid = tl.program_id(axis=0) + start = pid * CHUNK + h = tl.zeros((), dtype=tl.uint64) + for i in tl.static_range(0, CHUNK): + idx = start + i + m = idx < n_elems + v = tl.load(in_ptr + idx, mask=m, other=0).to(tl.uint64) + h += v + tl.store(out_ptr + pid, h) + + +def _as_uint32_words(t: torch.Tensor) -> torch.Tensor: + assert t.is_cuda, "Use .cuda() first" + tb = t.contiguous().view(torch.uint8) + nbytes = tb.numel() + pad = (4 - (nbytes & 3)) & 3 + if pad: + tb_p = torch.empty(nbytes + pad, dtype=torch.uint8, device=tb.device) + tb_p[:nbytes].copy_(tb) + tb_p[nbytes:].zero_() + tb = tb_p + return tb.view(torch.uint32) + + +def _final_splitmix64(x: int) -> int: + mask = (1 << 64) - 1 + x &= mask + x ^= x >> 30 + x = (x * 0xBF58476D1CE4E5B9) & mask + x ^= x >> 27 + x = (x * 0x94D049BB133111EB) & mask + x ^= x >> 31 + return x + + +@torch.inference_mode() +def gpu_tensor_hash( + tensor: torch.Tensor, + *, + seed: int = 0x243F6A88, + tile_words: int = 8192, + block_words: int = 256, + reduce_chunk: int = 1024, + num_warps: int = 4, + num_stages: int = 4, + use_cg: bool = True, +) -> int: + assert tensor.is_cuda, "Use .cuda() first" + u32 = _as_uint32_words(tensor) + n = u32.numel() + if n == 0: + return 0 + + grid1 = (triton.cdiv(n, tile_words),) + partials = torch.empty(grid1[0], dtype=torch.uint64, device=u32.device) + hash_tiles32_kernel_blocked[grid1]( + u32, + partials, + n, + seed1=seed & 0xFFFFFFFF, + seed2=((seed * 0x9E3779B1) ^ 0xDEADBEEF) & 0xFFFFFFFF, + FM_C1=FMIX32_C1, + FM_C2=FMIX32_C2, + POS_A=POS_C1, + POS_B=POS_C2, + TILE=tile_words, + BLOCK=block_words, + USE_CG=use_cg, + num_warps=num_warps, + num_stages=num_stages, + ) + + cur = partials + while cur.numel() > 1: + n_elems = cur.numel() + grid2 = (triton.cdiv(n_elems, reduce_chunk),) + nxt = torch.empty(grid2[0], dtype=torch.uint64, device=cur.device) + add_tree_reduce_u64_kernel[grid2](cur, nxt, n_elems, CHUNK=reduce_chunk) + cur = nxt + + return _final_splitmix64(int(cur.item())) diff --git a/python/sglang/srt/layers/parameter.py b/python/sglang/srt/layers/parameter.py new file mode 100644 index 000000000..1ea75d70c --- /dev/null +++ b/python/sglang/srt/layers/parameter.py @@ -0,0 +1,533 @@ +"""Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/parameter.py""" + +import logging +from fractions import Fraction +from typing import Callable, Optional, Union + +import torch +from torch.nn import Parameter + +from sglang.srt.utils import is_cpu + +__all__ = [ + "BasevLLMParameter", + "PackedvLLMParameter", + "PerTensorScaleParameter", + "ModelWeightParameter", + "ChannelQuantScaleParameter", + "GroupQuantScaleParameter", + "BlockQuantScaleParameter", + "PackedColumnParameter", + "RowvLLMParameter", +] + +logger = logging.getLogger(__name__) + +_is_cpu = is_cpu() + + +class BasevLLMParameter(Parameter): + """ + Base parameter for vLLM linear layers. Extends the torch.nn.parameter + by taking in a linear weight loader. Will copy the loaded weight + into the parameter when the provided weight loader is called. + """ + + def __new__(cls, data: torch.Tensor, **kwargs): + + return super().__new__(cls, data=data, requires_grad=False) + + def __init__(self, data: torch.Tensor, weight_loader: Callable): + """ + Initialize the BasevLLMParameter + + :param data: torch tensor with the parameter data + :param weight_loader: weight loader callable + + :returns: a torch.nn.parameter + """ + + self._weight_loader = weight_loader + + @property + def weight_loader(self): + return self._weight_loader + + def _assert_and_load(self, loaded_weight: torch.Tensor): + assert self.data.shape == loaded_weight.shape + self.data.copy_(loaded_weight) + + def load_column_parallel_weight(self, loaded_weight: torch.Tensor): + self._assert_and_load(loaded_weight) + + def load_row_parallel_weight(self, loaded_weight: torch.Tensor): + self._assert_and_load(loaded_weight) + + def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs): + self._assert_and_load(loaded_weight) + + def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs): + self._assert_and_load(loaded_weight) + + +class _ColumnvLLMParameter(BasevLLMParameter): + """ + Private class defining weight loading functionality + (load_merged_column_weight, load_qkv_weight) + for parameters being loaded into linear layers with column + parallelism. This includes QKV and MLP layers which are + not already fused on disk. Requires an output dimension + to be defined. Called within the weight loader of + each of the column parallel linear layers. + """ + + def __init__(self, output_dim: int, **kwargs): + self._output_dim = output_dim + super().__init__(**kwargs) + + @property + def output_dim(self): + return self._output_dim + + def load_column_parallel_weight( + self, + loaded_weight: torch.Tensor, + tp_rank: int, + use_presharded_weights: bool = False, + ): + if not use_presharded_weights: + shard_size = self.data.shape[self.output_dim] + + from sglang.srt.model_loader.weight_utils import ( + narrow_padded_param_and_loaded_weight, + ) + + if _is_cpu: + param_data, loaded_weight = narrow_padded_param_and_loaded_weight( + self.data, + loaded_weight, + 0, # param_data_start + tp_rank * shard_size, + self.output_dim, + shard_size, + ) + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + return + else: + loaded_weight = loaded_weight.narrow( + self.output_dim, tp_rank * shard_size, shard_size + ) + + assert self.data.shape == loaded_weight.shape + self.data.copy_(loaded_weight) + + def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs): + + shard_offset = kwargs.get("shard_offset") + shard_size = kwargs.get("shard_size") + tp_rank = kwargs.get("tp_rank") + use_presharded_weights = kwargs.get("use_presharded_weights") + if ( + isinstance(self, (PackedColumnParameter, PackedvLLMParameter)) + and self.packed_dim == self.output_dim + ): + shard_size, shard_offset = self.adjust_shard_indexes_for_packing( + shard_offset=shard_offset, shard_size=shard_size + ) + + param_data = self.data + + param_data = param_data.narrow(self.output_dim, shard_offset, shard_size) + + from sglang.srt.model_loader.weight_utils import ( + narrow_padded_param_and_loaded_weight, + ) + + if _is_cpu: + param_data, loaded_weight = narrow_padded_param_and_loaded_weight( + param_data, + loaded_weight, + 0, # param_data_start + tp_rank * shard_size, + self.output_dim, + shard_size, + not use_presharded_weights, + ) + else: + if not use_presharded_weights: + loaded_weight = loaded_weight.narrow( + self.output_dim, tp_rank * shard_size, shard_size + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def load_qkv_weight( + self, + loaded_weight: torch.Tensor, + tp_rank: int, + use_presharded_weights: bool = False, + **kwargs, + ): + + shard_offset = kwargs.get("shard_offset") + shard_size = kwargs.get("shard_size") + shard_id = kwargs.get("shard_id") + num_heads = kwargs.get("num_heads") + + if ( + isinstance(self, (PackedColumnParameter, PackedvLLMParameter)) + and self.output_dim == self.packed_dim + ): + shard_size, shard_offset = self.adjust_shard_indexes_for_packing( + shard_offset=shard_offset, shard_size=shard_size + ) + + param_data = self.data + shard_id = tp_rank if shard_id == "q" else tp_rank // num_heads + param_data = param_data.narrow(self.output_dim, shard_offset, shard_size) + + if _is_cpu: + from sglang.srt.model_loader.weight_utils import ( + narrow_padded_param_and_loaded_weight, + ) + + param_data, loaded_weight = narrow_padded_param_and_loaded_weight( + param_data, + loaded_weight, + 0, # param_data_start + shard_id * shard_size, + self.output_dim, + shard_size, + not use_presharded_weights, + ) + else: + if not use_presharded_weights: + loaded_weight = loaded_weight.narrow( + self.output_dim, shard_id * shard_size, shard_size + ) + + assert ( + param_data.shape == loaded_weight.shape + ), f"{param_data.shape=}, {loaded_weight.shape=}" + param_data.copy_(loaded_weight) + + +class RowvLLMParameter(BasevLLMParameter): + """ + Parameter class defining weight_loading functionality + (load_row_parallel_weight) for parameters being loaded + into linear layers with row parallel functionality. + Requires an input_dim to be defined. + """ + + def __init__(self, input_dim: int, **kwargs): + self._input_dim = input_dim + super().__init__(**kwargs) + + @property + def input_dim(self): + return self._input_dim + + def load_row_parallel_weight( + self, + loaded_weight: torch.Tensor, + tp_rank: int, + use_presharded_weights: bool = False, + ): + if not use_presharded_weights: + shard_size = self.data.shape[self.input_dim] + + from sglang.srt.model_loader.weight_utils import ( + narrow_padded_param_and_loaded_weight, + ) + + if _is_cpu: + param_data, loaded_weight = narrow_padded_param_and_loaded_weight( + self.data, + loaded_weight, + 0, # param_data_start + tp_rank * shard_size, + self.input_dim, + shard_size, + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + return + else: + loaded_weight = loaded_weight.narrow( + self.input_dim, tp_rank * shard_size, shard_size + ) + + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert self.data.shape == loaded_weight.shape + self.data.copy_(loaded_weight) + + +class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter): + """ + Parameter class for linear layer weights. Uses both column and + row parallelism. + """ + + pass + + +class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): + """ + Parameter class for weight scales loaded for weights with + grouped quantization. Uses both column and row parallelism. + """ + + pass + + +class ChannelQuantScaleParameter(_ColumnvLLMParameter): + """ + Parameter class for weight scales loaded for weights with + channel-wise quantization. Equivalent to _ColumnvLLMParameter. + """ + + pass + + +class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): + """ + Parameter class for weight scales loaded for weights with + block-wise quantization. Uses both column and row parallelism. + """ + + pass + + +class PerTensorScaleParameter(BasevLLMParameter): + """ + Parameter class for scales where the number of scales is + equivalent to the number of logical matrices in fused linear + layers (e.g. for QKV, there are 3 scales loaded from disk). + This is relevant to weights with per-tensor quantization. + Adds functionality to map the scalers to a shard during + weight loading. + + Note: additional parameter manipulation may be handled + for each quantization config specifically, within + process_weights_after_loading + """ + + def __init__(self, **kwargs): + self.qkv_idxs = {"q": 0, "k": 1, "v": 2} + super().__init__(**kwargs) + + def _shard_id_as_int(self, shard_id: Union[str, int]) -> int: + if isinstance(shard_id, int): + return shard_id + + # if not int, assume shard_id for qkv + # map to int and return + assert isinstance(shard_id, str) + assert shard_id in self.qkv_idxs + return self.qkv_idxs[shard_id] + + # For row parallel layers, no sharding needed + # load weight into parameter as is + def load_row_parallel_weight(self, *args, **kwargs): + kwargs.pop("tp_rank", None) + kwargs.pop("use_presharded_weights", None) + super().load_row_parallel_weight(*args, **kwargs) + + def load_merged_column_weight(self, *args, **kwargs): + self._load_into_shard_id(*args, **kwargs) + + def load_qkv_weight(self, *args, **kwargs): + self._load_into_shard_id(*args, **kwargs) + + def load_column_parallel_weight(self, *args, **kwargs): + kwargs.pop("tp_rank", None) + kwargs.pop("use_presharded_weights", None) + super().load_row_parallel_weight(*args, **kwargs) + + def _load_into_shard_id( + self, loaded_weight: torch.Tensor, shard_id: Union[str, int], **kwargs + ): + """ + Slice the parameter data based on the shard id for + loading. + """ + + param_data = self.data + shard_id = self._shard_id_as_int(shard_id) + + # AutoFP8 scales do not have a shape + # compressed-tensors scales do have a shape + if len(loaded_weight.shape) != 0: + assert loaded_weight.shape[0] == 1 + loaded_weight = loaded_weight[0] + + param_data = param_data[shard_id] + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class PackedColumnParameter(_ColumnvLLMParameter): + """ + Parameter for model parameters which are packed on disk + and support column parallelism only. See PackedvLLMParameter + for more details on the packed properties. + """ + + def __init__( + self, + packed_factor: Union[int, Fraction], + packed_dim: int, + marlin_tile_size: Optional[int] = None, + **kwargs, + ): + self._packed_factor = packed_factor + self._packed_dim = packed_dim + self._marlin_tile_size = marlin_tile_size + super().__init__(**kwargs) + + @property + def packed_dim(self): + return self._packed_dim + + @property + def packed_factor(self): + return self._packed_factor + + @property + def marlin_tile_size(self): + return self._marlin_tile_size + + def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): + return _adjust_shard_indexes_for_packing( + shard_size=shard_size, + shard_offset=shard_offset, + packed_factor=self.packed_factor, + marlin_tile_size=self.marlin_tile_size, + ) + + +class PackedvLLMParameter(ModelWeightParameter): + """ + Parameter for model weights which are packed on disk. + Example: GPTQ Marlin weights are int4 or int8, packed into int32. + Extends the ModelWeightParameter to take in the + packed factor, the packed dimension, and optionally, marlin + tile size for marlin kernels. Adjusts the shard_size and + shard_offset for fused linear layers model weight loading + by accounting for packing and optionally, marlin tile size. + """ + + def __init__( + self, + packed_factor: Union[int, Fraction], + packed_dim: int, + marlin_tile_size: Optional[int] = None, + **kwargs, + ): + self._packed_factor = packed_factor + self._packed_dim = packed_dim + self._marlin_tile_size = marlin_tile_size + super().__init__(**kwargs) + + @property + def packed_dim(self): + return self._packed_dim + + @property + def packed_factor(self): + return self._packed_factor + + @property + def marlin_tile_size(self): + return self._marlin_tile_size + + def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): + return _adjust_shard_indexes_for_packing( + shard_size=shard_size, + shard_offset=shard_offset, + packed_factor=self.packed_factor, + marlin_tile_size=self.marlin_tile_size, + ) + + +def permute_param_layout_( + param: BasevLLMParameter, input_dim: int, output_dim: int, **kwargs +) -> BasevLLMParameter: + """ + Permute a parameter's layout to the specified input and output dimensions, + useful for forcing the parameter into a known layout, for example, if I need + a packed (quantized) weight matrix to be in the layout + {input_dim = 0, output_dim = 1, packed_dim = 0} + then I can call: + permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) + to ensure x is in the correct layout (permuting it to the correct layout if + required, asserting if it cannot get it to the correct layout) + """ + + curr_input_dim = getattr(param, "input_dim", None) + curr_output_dim = getattr(param, "output_dim", None) + + if curr_input_dim is None or curr_output_dim is None: + assert param.data.dim() == 2, ( + "permute_param_layout_ only supports 2D parameters when either " + "input_dim or output_dim is not set" + ) + + # if one of the dimensions is not set, set it to the opposite of the other + # we can only do this since we asserted the parameter is 2D above + if curr_input_dim is None: + assert curr_output_dim is not None, "either input or output dim must be set" + curr_input_dim = (curr_output_dim + 1) % 2 + if curr_output_dim is None: + assert curr_input_dim is not None, "either input or output dim must be set" + curr_output_dim = (curr_input_dim + 1) % 2 + + # create permutation from the current layout to the layout with + # self.input_dim at input_dim and self.output_dim at output_dim preserving + # other dimensions + perm = [ + i for i in range(param.data.dim()) if i not in [curr_input_dim, curr_output_dim] + ] + perm.insert(input_dim, curr_input_dim) + perm.insert(output_dim, curr_output_dim) + + if "packed_dim" in kwargs: + assert ( + hasattr(param, "packed_dim") + and param.packed_dim == perm[kwargs["packed_dim"]] + ), "permute_param_layout_ currently doesn't support repacking" + + param.data = param.data.permute(*perm) + if hasattr(param, "_input_dim"): + param._input_dim = input_dim + if hasattr(param, "_output_dim"): + param._output_dim = output_dim + if "packed_dim" in kwargs and hasattr(param, "_packed_dim"): + param._packed_dim = kwargs["packed_dim"] + + return param + + +def _adjust_shard_indexes_for_marlin(shard_size, shard_offset, marlin_tile_size): + return shard_size * marlin_tile_size, shard_offset * marlin_tile_size + + +def _adjust_shard_indexes_for_packing( + shard_size, shard_offset, packed_factor, marlin_tile_size +): + shard_size = shard_size // packed_factor + shard_offset = shard_offset // packed_factor + if marlin_tile_size is not None: + return _adjust_shard_indexes_for_marlin( + shard_size=shard_size, + shard_offset=shard_offset, + marlin_tile_size=marlin_tile_size, + ) + return shard_size, shard_offset diff --git a/python/sglang/srt/layers/pooler.py b/python/sglang/srt/layers/pooler.py new file mode 100644 index 000000000..26bc5899e --- /dev/null +++ b/python/sglang/srt/layers/pooler.py @@ -0,0 +1,112 @@ +# adapted from +# https://github.com/vllm-project/vllm/blob/82a1b1a82b1fbb454c82a9ef95730b929c9b270c/vllm/model_executor/layers/pooler.py + +from dataclasses import dataclass +from enum import IntEnum +from typing import Optional + +import torch +import torch.nn as nn +from transformers import PretrainedConfig + +from sglang.srt.layers.activation import get_cross_encoder_activation_function +from sglang.srt.model_executor.model_runner import ForwardBatch + + +class PoolingType(IntEnum): + LAST = 0 + CLS = 1 + + +@dataclass +class EmbeddingPoolerOutput: + embeddings: torch.Tensor + + +class Pooler(nn.Module): + """A layer that pools specific information from hidden states. + This layer does the following: + 1. Extracts specific tokens or aggregates data based on pooling method. + 2. Normalizes output if specified. + 3. Returns structured results as `PoolerOutput`. + Attributes: + pooling_type: The type of pooling to use (LAST, AVERAGE, MAX). + normalize: Whether to normalize the pooled data. + """ + + def __init__(self, pooling_type: PoolingType, normalize: bool): + super().__init__() + self.pooling_type = pooling_type + self.normalize = normalize + + def forward( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ) -> EmbeddingPoolerOutput: + if self.pooling_type == PoolingType.LAST: + last_token_indices = torch.cumsum(forward_batch.extend_seq_lens, dim=0) - 1 + pooled_data = hidden_states[last_token_indices] + elif self.pooling_type == PoolingType.CLS: + prompt_lens = forward_batch.extend_seq_lens + first_token_flat_indices = torch.zeros_like(prompt_lens) + first_token_flat_indices[1:] += torch.cumsum(prompt_lens, dim=0)[:-1] + pooled_data = hidden_states[first_token_flat_indices] + else: + raise ValueError(f"Invalid pooling type: {self.pooling_type}") + + if self.normalize: + pooled_data = nn.functional.normalize(pooled_data, p=2, dim=1) + + return EmbeddingPoolerOutput(embeddings=pooled_data) + + +class CrossEncodingPooler(nn.Module): + """A layer that pools specific information from hidden states. + + This layer does the following: + 1. Extracts specific tokens or aggregates data based on pooling method. + 2. Normalizes output if specified. + 3. Returns structured results as `EmbeddingPoolerOutput`. + """ + + def __init__( + self, + config: PretrainedConfig, + classifier: nn.Module, + pooler: Optional[nn.Module] = None, + ): + super().__init__() + self.classifier = classifier + self.pooler = pooler + self.default_activation_function = get_cross_encoder_activation_function(config) + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> EmbeddingPoolerOutput: + """Pools sentence pair scores from the hidden_states.""" + + prompt_lens = forward_batch.extend_seq_lens + + offset = 0 + pooled_data_lst = [] + for prompt_len in prompt_lens: + pooled_data_i = hidden_states[offset : offset + prompt_len] + + if self.pooler is not None: + final_shape_tensor = self.pooler(pooled_data_i, forward_batch) + else: + final_shape_tensor = self.classifier(pooled_data_i) + + pooled_data_lst.append(final_shape_tensor) + offset += prompt_len + + pooled_output = torch.stack(pooled_data_lst) + + if self.pooler is not None: + # apply classifier once on the full batch if possible + pooled_output = self.classifier(pooled_output) + + scores = self.default_activation_function(pooled_output).squeeze(-1) + + return EmbeddingPoolerOutput(embeddings=scores) diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py new file mode 100644 index 000000000..ff3c2b148 --- /dev/null +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -0,0 +1,224 @@ +# Adapted from https://raw.githubusercontent.com/vllm-project/vllm/v0.5.5/vllm/model_executor/layers/quantization/__init__.py +from __future__ import annotations + +import builtins +import inspect +from typing import TYPE_CHECKING, Dict, Optional, Type + +import torch + +try: + from vllm.model_executor.layers.quantization.aqlm import AQLMConfig + from vllm.model_executor.layers.quantization.bitsandbytes import BitsAndBytesConfig + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( + CompressedTensorsW8A8Fp8MoEMethod, + CompressedTensorsWNA16MoEMethod, + ) + from vllm.model_executor.layers.quantization.deepspeedfp import DeepSpeedFPConfig + from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config + from vllm.model_executor.layers.quantization.gguf import GGUFConfig + from vllm.model_executor.layers.quantization.gptq_marlin_24 import ( + GPTQMarlin24Config, + ) + from vllm.model_executor.layers.quantization.marlin import MarlinConfig + from vllm.model_executor.layers.quantization.qqq import QQQConfig + from vllm.model_executor.layers.quantization.tpu_int8 import Int8TpuConfig + + VLLM_AVAILABLE = True +except ImportError as e: + VLLM_AVAILABLE = False + VLLM_IMPORT_ERROR = e + + # Define empty classes as placeholders when vllm is not available + class DummyConfig: + def override_quantization_method(self, *args, **kwargs): + return None + + AQLMConfig = BitsAndBytesConfig = CompressedTensorsConfig = DeepSpeedFPConfig = ( + ExpertsInt8Config + ) = GGUFConfig = GPTQMarlin24Config = MarlinConfig = QQQConfig = Int8TpuConfig = ( + DummyConfig + ) + + +from sglang.srt.layers.quantization.awq import AWQConfig, AWQMarlinConfig +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.blockwise_int8 import BlockInt8Config +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, +) +from sglang.srt.layers.quantization.fp8 import Fp8Config +from sglang.srt.layers.quantization.fpgemm_fp8 import FBGEMMFp8Config +from sglang.srt.layers.quantization.gptq import GPTQConfig, GPTQMarlinConfig +from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptFp4Config, + ModelOptFp8Config, +) +from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config +from sglang.srt.layers.quantization.mxfp4 import Mxfp4Config +from sglang.srt.layers.quantization.petit import PetitNvFp4Config +from sglang.srt.layers.quantization.qoq import QoQConfig +from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config +from sglang.srt.layers.quantization.w8a8_fp8 import W8A8Fp8Config +from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config +from sglang.srt.utils import is_cuda, is_hip, mxfp_supported + +_is_mxfp_supported = mxfp_supported() + +if TYPE_CHECKING: + from sglang.srt.layers.moe.topk import TopKOutput + +# Base quantization methods that don't depend on vllm +BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = { + "fp8": Fp8Config, + "blockwise_int8": BlockInt8Config, + "modelopt": ModelOptFp8Config, + "modelopt_fp4": ModelOptFp4Config, + "w8a8_int8": W8A8Int8Config, + "w8a8_fp8": W8A8Fp8Config, + "awq": AWQConfig, + "awq_marlin": AWQMarlinConfig, + "gptq": GPTQConfig, + "gptq_marlin": GPTQMarlinConfig, + "moe_wna16": MoeWNA16Config, + "compressed-tensors": CompressedTensorsConfig, + "qoq": QoQConfig, + "w4afp8": W4AFp8Config, + "petit_nvfp4": PetitNvFp4Config, + "fbgemm_fp8": FBGEMMFp8Config, +} + + +if is_cuda(): + BASE_QUANTIZATION_METHODS.update( + { + "quark": Mxfp4Config, + "mxfp4": Mxfp4Config, + } + ) +elif _is_mxfp_supported and is_hip(): + from sglang.srt.layers.quantization.quark.quark import QuarkConfig + + BASE_QUANTIZATION_METHODS.update( + { + "quark": QuarkConfig, + "mxfp4": Mxfp4Config, + } + ) +# VLLM-dependent quantization methods +VLLM_QUANTIZATION_METHODS = { + "aqlm": AQLMConfig, + "deepspeedfp": DeepSpeedFPConfig, + "tpu_int8": Int8TpuConfig, + "marlin": MarlinConfig, + "gguf": GGUFConfig, + "gptq_marlin_24": GPTQMarlin24Config, + "bitsandbytes": BitsAndBytesConfig, + "qqq": QQQConfig, + "experts_int8": ExpertsInt8Config, +} + +QUANTIZATION_METHODS = {**BASE_QUANTIZATION_METHODS, **VLLM_QUANTIZATION_METHODS} + + +def get_quantization_config(quantization: str) -> Type[QuantizationConfig]: + if quantization not in QUANTIZATION_METHODS: + raise ValueError( + f"Invalid quantization method: {quantization}. " + f"Available methods: {list(QUANTIZATION_METHODS.keys())}" + ) + if quantization in VLLM_QUANTIZATION_METHODS and not VLLM_AVAILABLE: + raise ValueError( + f"{quantization} quantization requires some operators from vllm. " + f"Please install vllm by `pip install vllm==0.9.0.1`\n" + f"Import error: {VLLM_IMPORT_ERROR}" + ) + + return QUANTIZATION_METHODS[quantization] + + +original_isinstance = builtins.isinstance + + +def monkey_patch_isinstance_for_vllm_base_layer(reverse: bool = False): + """ + Patch isinstance so that the `get_quant_method` in vllm's QuantizationConfig + can recognize sglang layers + """ + if not VLLM_AVAILABLE: + return + + if reverse: + builtins.isinstance = original_isinstance + return + + from vllm.model_executor.layers.fused_moe import FusedMoE + from vllm.model_executor.layers.linear import LinearBase + from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, + ) + + from sglang.srt.layers.linear import LinearBase as PatchedLinearBase + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE as PatchedFusedMoE + from sglang.srt.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding as PatchedVocabParallelEmbedding, + ) + + def patched_isinstance(obj, classinfo): + if classinfo is LinearBase: + return original_isinstance(obj, PatchedLinearBase) + if classinfo is FusedMoE: + return original_isinstance(obj, PatchedFusedMoE) + if classinfo is VocabParallelEmbedding: + return original_isinstance(obj, PatchedVocabParallelEmbedding) + return original_isinstance(obj, classinfo) + + builtins.isinstance = patched_isinstance + + +def monkey_patch_moe_apply(class_obj: "FusedMoEMethodBase"): + """ + Monkey patch the apply function of vllm's FusedMoEMethodBase. + Convert sglang arguments to vllm arguments. + """ + original_apply = class_obj.apply + sig = inspect.signature(original_apply) + param_names = list(sig.parameters.keys()) + has_correction_bias = "e_score_correction_bias" in param_names + + def new_apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + topk_output: TopKOutput, + *, + activation: str = "silu", + apply_router_weight_on_input: bool = False, + inplace: bool = True, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + ): + assert activation == "silu" + assert inplace and not no_combine + + kwargs = { + "self": self, + "layer": layer, + "x": x, + "topk_output": topk_output, + } + return original_apply(**kwargs) + + setattr(class_obj, "apply", new_apply) + + +def monkey_patch_quant_configs(): + """Apply all monkey patches in one place.""" + + monkey_patch_moe_apply(CompressedTensorsW8A8Fp8MoEMethod) + monkey_patch_moe_apply(CompressedTensorsWNA16MoEMethod) + + +# Only apply monkey patches if vllm is available +if VLLM_AVAILABLE: + monkey_patch_quant_configs() diff --git a/python/sglang/srt/layers/quantization/awq.py b/python/sglang/srt/layers/quantization/awq.py new file mode 100644 index 000000000..9cba60c2b --- /dev/null +++ b/python/sglang/srt/layers/quantization/awq.py @@ -0,0 +1,782 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import logging +import warnings +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional + +import torch + +from sglang.srt.layers.linear import LinearBase, set_weight_attrs +from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.marlin_utils import ( + apply_awq_marlin_linear, + awq_to_marlin_zero_points, + check_marlin_supported, + check_marlin_supports_layer, + check_moe_marlin_supports_layer, + marlin_make_empty_g_idx, + marlin_make_workspace, + marlin_moe_permute_scales, + marlin_permute_scales, + moe_awq_to_marlin_zero_points, + verify_marlin_supported, + verify_marlin_supports_shape, +) +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter + +if TYPE_CHECKING: + from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig + from sglang.srt.layers.moe.token_dispatcher import ( + StandardDispatchOutput, + CombineInput, + ) + +from sglang.srt.utils import is_cuda, is_hip + +_is_cuda = is_cuda() +_is_hip = is_hip() +if _is_cuda: + from sgl_kernel import ( + awq_dequantize, + awq_marlin_moe_repack, + awq_marlin_repack, + fused_marlin_moe, + ) + + +elif _is_hip: + from sglang.srt.layers.quantization.awq_triton import ( + awq_dequantize_triton as awq_dequantize, + ) + + warnings.warn(f"HIP does not support fused_marlin_moe currently.") +else: + warnings.warn(f"Only CUDA and HIP support AWQ currently.") + +logger = logging.getLogger(__name__) + + +ScalarType, scalar_types = get_scalar_types() + + +def is_layer_skipped_awq(prefix: str, modules_to_not_convert: List[str]): + return any(module_name in prefix for module_name in modules_to_not_convert) + + +class AWQConfig(QuantizationConfig): + """Config class for AWQ. + + Reference: https://arxiv.org/abs/2306.00978 + """ + + def __init__( + self, + weight_bits: int, + group_size: int, + zero_point: bool, + modules_to_not_convert: Optional[List[str]] = None, + ) -> None: + super().__init__() + self.weight_bits = weight_bits + self.group_size = group_size + self.zero_point = zero_point + self.modules_to_not_convert = modules_to_not_convert or [] + + if self.weight_bits != 4: + raise ValueError( + "Currently, only 4-bit weight quantization is supported for " + f"AWQ, but got {self.weight_bits} bits." + ) + self.pack_factor = 32 // self.weight_bits + + def __repr__(self) -> str: + return ( + f"AWQConfig(weight_bits={self.weight_bits}, " + f"group_size={self.group_size}, " + f"zero_point={self.zero_point}, " + f"modules_to_not_convert={self.modules_to_not_convert})" + ) + + def get_scaled_act_names(self) -> List[str]: + return [] + + def get_name(self) -> str: + return "awq" + + def get_supported_act_dtypes(self) -> List[torch.dtype]: + return [torch.half] + + @classmethod + def get_min_capability(cls) -> int: + # The AWQ kernel only supports Turing or newer GPUs. + return 75 + + @staticmethod + def get_config_filenames() -> List[str]: + return [ + "quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq + # E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq + "quantize_config.json", + ] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> AWQConfig: + weight_bits = cls.get_from_keys(config, ["w_bit", "bits"]) + group_size = cls.get_from_keys(config, ["q_group_size", "group_size"]) + zero_point = cls.get_from_keys(config, ["zero_point"]) + modules_to_not_convert = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + return cls(weight_bits, group_size, zero_point, modules_to_not_convert) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[LinearMethodBase]: + from sglang.srt.layers.linear import LinearBase + + if isinstance(layer, LinearBase): + if is_layer_skipped_awq(prefix, self.modules_to_not_convert): + return UnquantizedLinearMethod() + return AWQLinearMethod(self) + return None + + +class AWQMarlinConfig(QuantizationConfig): + """Config class for AWQ Marlin""" + + # num_bits -> type + TYPE_MAP = { + 4: scalar_types.uint4, + 8: scalar_types.uint8, + } + + def __init__( + self, + weight_bits: int, + group_size: int, + zero_point: bool, + lm_head_quantized: bool, + modules_to_not_convert: Optional[list[str]], + full_config: dict[str, Any], + ) -> None: + super().__init__() + self.pack_factor = 32 // weight_bits # packed into int32 + self.group_size = group_size + self.zero_point = zero_point + self.lm_head_quantized = lm_head_quantized + self.weight_bits = weight_bits + self.modules_to_not_convert = modules_to_not_convert or [] + self.full_config = full_config + + if self.weight_bits not in self.TYPE_MAP: + raise ValueError( + f"Unsupported num_bits = {self.weight_bits}. " + f"Supported num_bits = {self.TYPE_MAP.keys()}" + ) + + self.quant_type = self.TYPE_MAP[self.weight_bits] + + verify_marlin_supported( + self.quant_type, group_size=self.group_size, has_zp=self.zero_point + ) + + def __repr__(self) -> str: + return ( + f"AWQMarlinConfig(quant_type={self.quant_type}, " + f"group_size={self.group_size}, " + f"zero_point={self.zero_point}, " + f"lm_head_quantized={self.lm_head_quantized}, " + f"modules_to_not_convert={self.modules_to_not_convert})" + ) + + def get_scaled_act_names(self) -> List[str]: + return [] + + @classmethod + def get_name(cls) -> str: + return "awq_marlin" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.half, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return ["quantize_config.json"] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> AWQMarlinConfig: + weight_bits = cls.get_from_keys(config, ["bits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + zero_point = cls.get_from_keys(config, ["zero_point"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + modules_to_not_convert = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + return cls( + weight_bits, + group_size, + zero_point, + lm_head_quantized, + modules_to_not_convert, + config, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg) + is_valid_user_quant = ( + user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin" + ) + + if can_convert and is_valid_user_quant: + msg = ( + "The model is convertible to {} during runtime." + " Using {} kernel.".format(cls.get_name(), cls.get_name()) + ) + logger.info(msg) + return cls.get_name() + + if can_convert and user_quant == "awq": + logger.info( + "Detected that the model can run with awq_marlin" + ", however you specified quantization=awq explicitly," + " so forcing awq. Use quantization=awq_marlin for" + " faster inference" + ) + return None + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + + if isinstance(layer, LinearBase) or ( + isinstance(layer, ParallelLMHead) and self.lm_head_quantized + ): + if is_layer_skipped_awq(prefix, self.modules_to_not_convert): + return UnquantizedLinearMethod() + # Check if the layer is supported by AWQMarlin. + if not check_marlin_supports_layer(layer, self.group_size): + logger.warning_once( + "Layer '%s' is not supported by AWQMarlin. Falling back to unoptimized AWQ kernels.", # noqa: E501 + prefix, + ) + return AWQConfig.from_config(self.full_config).get_quant_method( + layer, prefix + ) + return AWQMarlinLinearMethod(self) + elif isinstance(layer, FusedMoE): + from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config + + if not check_moe_marlin_supports_layer(layer, self.group_size): + logger.warning_once( + f"Layer '{prefix}' is not supported by AWQMoeMarlin. " + "Falling back to Moe WNA16 kernels." + ) + return MoeWNA16Config.from_config(self.full_config).get_quant_method( + layer, prefix + ) + return AWQMoEMethod(self) + return None + + @classmethod + def is_awq_marlin_compatible(cls, quant_config: dict[str, Any]): + # Extract data from quant config. + quant_method = quant_config.get("quant_method", "").lower() + num_bits = quant_config.get("bits") + group_size = quant_config.get("group_size") + zero_point = quant_config.get("zero_point") + + if not _is_cuda: + return False + + if quant_method != "awq": + return False + + # If we cannot find the info needed in the config, cannot convert. + if num_bits is None or group_size is None or zero_point is None: + return False + + if num_bits not in cls.TYPE_MAP: + return False + + return check_marlin_supported( + quant_type=cls.TYPE_MAP[num_bits], group_size=group_size, has_zp=zero_point + ) + + +class AWQLinearMethod(LinearMethodBase): + """Linear method for AWQ. + + Args: + quant_config: The AWQ quantization config. + """ + + def __init__(self, quant_config: AWQConfig): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + if input_size_per_partition % self.quant_config.group_size != 0: + raise ValueError( + "The input size is not aligned with the quantized " + "weight shape. This can be caused by too large " + "tensor parallel size." + ) + + output_size_per_partition = sum(output_partition_sizes) + if output_size_per_partition % self.quant_config.pack_factor != 0: + raise ValueError( + "The output size is not aligned with the quantized " + "weight shape. This can be caused by too large " + "tensor parallel size." + ) + + weight_loader = extra_weight_attrs.get("weight_loader") + qweight = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + qzeros = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition // self.quant_config.group_size, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + scales = GroupQuantScaleParameter( + data=torch.empty( + input_size_per_partition // self.quant_config.group_size, + output_size_per_partition, + dtype=params_dtype, + ), + input_dim=0, + output_dim=1, + weight_loader=weight_loader, + ) + + layer.register_parameter("qweight", qweight) + layer.register_parameter("qzeros", qzeros) + layer.register_parameter("scales", scales) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) + layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) + layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + qweight = layer.qweight + scales = layer.scales + qzeros = layer.qzeros + pack_factor = self.quant_config.pack_factor + out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,) + reshaped_x = x.reshape(-1, x.shape[-1]) + out = awq_dequantize(qweight, scales, qzeros) + out = torch.matmul(reshaped_x, out) + + if bias is not None: + out.add_(bias) + return out.reshape(out_shape) + + +class AWQMarlinLinearMethod(LinearMethodBase): + """Linear method for AWQ Marlin. + + Args: + quant_config: The AWQ Marlin quantization config. + """ + + def __init__(self, quant_config: AWQMarlinConfig) -> None: + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del output_size + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + # Normalize group_size + if self.quant_config.group_size != -1: + group_size = self.quant_config.group_size + else: + group_size = input_size + + verify_marlin_supports_shape( + output_size_per_partition=output_size_per_partition, + input_size_per_partition=input_size_per_partition, + input_size=input_size, + group_size=group_size, + ) + + qweight = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + num_groups = input_size_per_partition // group_size + + qzeros = PackedvLLMParameter( + data=torch.empty( + num_groups, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + scales = GroupQuantScaleParameter( + data=torch.empty( + num_groups, + output_size_per_partition, + dtype=params_dtype, + ), + input_dim=0, + output_dim=1, + weight_loader=weight_loader, + ) + + layer.register_parameter("qweight", qweight) + layer.register_parameter("qzeros", qzeros) + layer.register_parameter("scales", scales) + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.num_groups = num_groups + + # TODO: Update this docs + # Checkpoints are serialized in AutoAWQ format, which is different from the + # marlin format. This function is called after the weights are loaded. + # Here, we handle the repacking + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + device = layer.qweight.device + layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) + layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) + layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) + + # Allocate marlin workspace + layer.workspace = marlin_make_workspace(device) + + # Repack weights from AWQ format to marlin format. + marlin_qweight = awq_marlin_repack( + layer.qweight, + size_k=layer.input_size_per_partition, + size_n=layer.output_size_per_partition, + num_bits=self.quant_config.quant_type.size_bits, + ) + replace_parameter(layer, "qweight", marlin_qweight) + + # Permute scales from AWQ format to marlin format. + marlin_scales = marlin_permute_scales( + layer.scales, + size_k=layer.input_size_per_partition, + size_n=layer.output_size_per_partition, + group_size=self.quant_config.group_size, + ) + replace_parameter(layer, "scales", marlin_scales) + + # Permute zero-points from AWQ format to marlin format. + marlin_zp = awq_to_marlin_zero_points( + layer.qzeros, + size_k=layer.num_groups, + size_n=layer.output_size_per_partition, + num_bits=self.quant_config.quant_type.size_bits, + ) + replace_parameter(layer, "qzeros", marlin_zp) + + # Not-used + layer.g_idx = marlin_make_empty_g_idx(device) + layer.g_idx_sort_indices = marlin_make_empty_g_idx(device) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_awq_marlin_linear( + input=x, + weight=layer.qweight, + weight_scale=layer.scales, + weight_zp=layer.qzeros, + g_idx=layer.g_idx, + g_idx_sort_indices=layer.g_idx_sort_indices, + workspace=layer.workspace, + quant_type=self.quant_config.quant_type, + output_size_per_partition=layer.output_size_per_partition, + input_size_per_partition=layer.input_size_per_partition, + bias=bias, + ) + + +class AWQMoEMethod(FusedMoEMethodBase): + + def __init__(self, quant_config: AWQMarlinConfig): + self.quant_config = quant_config + if self.quant_config.weight_bits != 4: + raise ValueError("AWQMoEMethod only supports 4bit now.") + self.quant_type = scalar_types.uint4 + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + # Delay the import to avoid circular dependency + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + extra_weight_attrs.update( + { + "is_transposed": True, + "quant_method": FusedMoeWeightScaleSupported.GROUP.value, + } + ) + + w13_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + 2 * intermediate_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qweight", w13_qweight) + set_weight_attrs(w13_qweight, extra_weight_attrs) + + w2_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + hidden_size // self.quant_config.pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qweight", w2_qweight) + set_weight_attrs(w2_qweight, extra_weight_attrs) + + num_groups_w13 = hidden_size // self.quant_config.group_size + num_groups_w2 = intermediate_size_per_partition // self.quant_config.group_size + + # WEIGHT_SCALES + # Allocate 2 scales for w1 and w3 respectively. + w13_scales = torch.nn.Parameter( + torch.empty( + num_experts, + num_groups_w13, + intermediate_size_per_partition * 2, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_scales", w13_scales) + set_weight_attrs(w13_scales, extra_weight_attrs) + + w2_scales = torch.nn.Parameter( + torch.empty(num_experts, num_groups_w2, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_scales", w2_scales) + set_weight_attrs(w2_scales, extra_weight_attrs) + + # WEIGHT_ZERO_POINT + # Allocate 2 zero points for w1 and w3 respectively. + w13_qzeros = torch.nn.Parameter( + torch.empty( + num_experts, + num_groups_w13, + 2 * intermediate_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qzeros", w13_qzeros) + set_weight_attrs(w13_qzeros, extra_weight_attrs) + + w2_qzeros = torch.nn.Parameter( + torch.empty( + num_experts, + num_groups_w2, + hidden_size // self.quant_config.pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qzeros", w2_qzeros) + set_weight_attrs(w2_qzeros, extra_weight_attrs) + + device = layer.w13_qweight.device + layer.workspace = marlin_make_workspace(device, 4) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + num_experts = layer.w13_qweight.shape[0] + device = layer.w13_qweight.device + + layer.w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + + marlin_w13_qweight = awq_marlin_moe_repack( + layer.w13_qweight, + layer.w13_g_idx_sort_indices, + size_k=layer.w13_qweight.shape[1], + size_n=layer.w13_qweight.shape[2] * self.quant_config.pack_factor, + num_bits=self.quant_config.weight_bits, + ) + replace_parameter(layer, "w13_qweight", marlin_w13_qweight) + + marlin_w2_qweight = awq_marlin_moe_repack( + layer.w2_qweight, + layer.w2_g_idx_sort_indices, + size_k=layer.w2_qweight.shape[1], + size_n=layer.w2_qweight.shape[2] * self.quant_config.pack_factor, + num_bits=self.quant_config.weight_bits, + ) + replace_parameter(layer, "w2_qweight", marlin_w2_qweight) + + # hidden_size->intermediate_size + marlin_w13_scales = marlin_moe_permute_scales( + s=layer.w13_scales, + size_k=layer.intermediate_size_per_partition, + size_n=layer.w13_scales.shape[2], + group_size=self.quant_config.group_size, + ) + + replace_parameter(layer, "w13_scales", marlin_w13_scales) + + marlin_w2_scales = marlin_moe_permute_scales( + s=layer.w2_scales, + size_k=layer.intermediate_size_per_partition, + size_n=layer.w2_scales.shape[2], + group_size=self.quant_config.group_size, + ) + replace_parameter(layer, "w2_scales", marlin_w2_scales) + + marlin_w13_zp = moe_awq_to_marlin_zero_points( + layer.w13_qzeros, + size_k=layer.w13_qzeros.shape[1], + size_n=layer.w13_qzeros.shape[2] * self.quant_config.pack_factor, + num_bits=self.quant_config.weight_bits, + ) + replace_parameter(layer, "w13_qzeros", marlin_w13_zp) + + marlin_w2_zp = moe_awq_to_marlin_zero_points( + layer.w2_qzeros, + size_k=layer.w2_qzeros.shape[1], + size_n=layer.w2_qzeros.shape[2] * self.quant_config.pack_factor, + num_bits=self.quant_config.weight_bits, + ) + replace_parameter(layer, "w2_qzeros", marlin_w2_zp) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + assert ( + self.moe_runner_config.activation == "silu" + ), "Only SiLU activation is supported." + + # The input must currently be float16 + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + orig_dtype = x.dtype + x = x.half() + + topk_weights, topk_ids, router_logits = topk_output + + output = fused_marlin_moe( + x, + layer.w13_qweight, + layer.w2_qweight, + layer.w13_scales, + layer.w2_scales, + router_logits, + topk_weights, + topk_ids, + sort_indices1=layer.w13_g_idx_sort_indices, + sort_indices2=layer.w2_g_idx_sort_indices, + w1_zeros=layer.w13_qzeros, + w2_zeros=layer.w2_qzeros, + num_bits=self.quant_config.weight_bits, + ).to(orig_dtype) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/awq_triton.py b/python/sglang/srt/layers/quantization/awq_triton.py new file mode 100644 index 000000000..13352efdb --- /dev/null +++ b/python/sglang/srt/layers/quantization/awq_triton.py @@ -0,0 +1,339 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/awq_triton.py + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import triton +import triton.language as tl + +AWQ_TRITON_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] + + +@triton.jit +def awq_dequantize_kernel( + qweight_ptr, # quantized matrix + scales_ptr, # scales, per group + zeros_ptr, # zeros, per group + group_size, # Should always be one of the supported group sizes + result_ptr, # Output matrix + num_cols, # input num cols in qweight + num_rows, # input num rows in qweight + BLOCK_SIZE_X: tl.constexpr, + BLOCK_SIZE_Y: tl.constexpr, +): + # Setup the pids. + pid_x = tl.program_id(axis=0) + pid_y = tl.program_id(axis=1) + + # Compute offsets and masks for qweight_ptr. + offsets_y = pid_y * BLOCK_SIZE_Y + tl.arange(0, BLOCK_SIZE_Y) + offsets_x = pid_x * BLOCK_SIZE_X + tl.arange(0, BLOCK_SIZE_X) + offsets = num_cols * offsets_y[:, None] + offsets_x[None, :] + + masks_y = offsets_y < num_rows + masks_x = offsets_x < num_cols + + masks = masks_y[:, None] & masks_x[None, :] + + # Compute offsets and masks for result output ptr. + result_offsets_y = pid_y * BLOCK_SIZE_Y + tl.arange(0, BLOCK_SIZE_Y) + result_offsets_x = pid_x * BLOCK_SIZE_X * 8 + tl.arange(0, BLOCK_SIZE_X * 8) + result_offsets = ( + 8 * num_cols * result_offsets_y[:, None] + result_offsets_x[None, :] + ) + + result_masks_y = result_offsets_y < num_rows + result_masks_x = result_offsets_x < num_cols * 8 + result_masks = result_masks_y[:, None] & result_masks_x[None, :] + + # Load the weights. + iweights = tl.load(qweight_ptr + offsets, masks, 0.0) + iweights = tl.interleave(iweights, iweights) + iweights = tl.interleave(iweights, iweights) + iweights = tl.interleave(iweights, iweights) + + # Create reverse AWQ order as tensor: [0, 4, 1, 5, 2, 6, 3, 7] + # that will map given indices to the correct order. + reverse_awq_order_tensor = ( + (tl.arange(0, 2) * 4)[None, :] + tl.arange(0, 4)[:, None] + ).reshape(8) + + # Use this to compute a set of shifts that can be used to unpack and + # reorder the values in iweights and zeros. + shifts = reverse_awq_order_tensor * 4 + shifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_Y * BLOCK_SIZE_X, 8)) + shifts = tl.reshape(shifts, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8)) + + # Unpack and reorder: shift out the correct 4-bit value and mask. + iweights = (iweights >> shifts) & 0xF + + # Compute zero offsets and masks. + zero_offsets_y = pid_y * BLOCK_SIZE_Y // group_size + tl.arange(0, 1) + zero_offsets_x = pid_x * BLOCK_SIZE_X + tl.arange(0, BLOCK_SIZE_X) + zero_offsets = num_cols * zero_offsets_y[:, None] + zero_offsets_x[None, :] + + zero_masks_y = zero_offsets_y < num_rows // group_size + zero_masks_x = zero_offsets_x < num_cols + zero_masks = zero_masks_y[:, None] & zero_masks_x[None, :] + + # Load the zeros. + zeros = tl.load(zeros_ptr + zero_offsets, zero_masks, 0.0) + zeros = tl.interleave(zeros, zeros) + zeros = tl.interleave(zeros, zeros) + zeros = tl.interleave(zeros, zeros) + zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8)) + + # Unpack and reorder: shift out the correct 4-bit value and mask. + zeros = (zeros >> shifts) & 0xF + + # Compute scale offsets and masks. + scale_offsets_y = pid_y * BLOCK_SIZE_Y // group_size + tl.arange(0, 1) + scale_offsets_x = pid_x * BLOCK_SIZE_X * 8 + tl.arange(0, BLOCK_SIZE_X * 8) + scale_offsets = num_cols * 8 * scale_offsets_y[:, None] + scale_offsets_x[None, :] + scale_masks_y = scale_offsets_y < num_rows // group_size + scale_masks_x = scale_offsets_x < num_cols * 8 + scale_masks = scale_masks_y[:, None] & scale_masks_x[None, :] + + # Load the scales. + scales = tl.load(scales_ptr + scale_offsets, scale_masks, 0.0) + scales = tl.broadcast_to(scales, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8)) + + # Dequantize. + iweights = (iweights - zeros) * scales + iweights = iweights.to(result_ptr.type.element_ty) + + # Finally, store. + tl.store(result_ptr + result_offsets, iweights, result_masks) + + +@triton.jit +def awq_gemm_kernel( + a_ptr, + b_ptr, + c_ptr, + zeros_ptr, + scales_ptr, + M, + N, + K, + group_size, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SPLIT_K: tl.constexpr, +): + pid = tl.program_id(axis=0) + pid_z = tl.program_id(1) + + # NOTE: This doesn't work in TRITON_INTERPRET=1 mode. Use below instead. + # num_pid_n = (N + BLOCK_SIZE_N - 1) // BLOCK_SIZE_N + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + accumulator_dtype = c_ptr.type.element_ty + + # NOTE: This doesn't work in TRITON_INTERPRET=1 mode. Use below instead. + # accumulator = tl.arange(0, BLOCK_SIZE_N) + # accumulator = tl.broadcast_to(accumulator[None, :], + # (BLOCK_SIZE_M, BLOCK_SIZE_N)) + # accumulator = accumulator & 0x0 + # accumulator = accumulator.to(accumulator_dtype) + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype) + + # Create reverse AWQ order as tensor: [0, 4, 1, 5, 2, 6, 3, 7] + # that will map given indices to the correct order. + reverse_awq_order_tensor = ( + (tl.arange(0, 2) * 4)[None, :] + tl.arange(0, 4)[:, None] + ).reshape(8) + + # Create the necessary shifts to use to unpack. + shifts = reverse_awq_order_tensor * 4 + shifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_K * (BLOCK_SIZE_N // 8), 8)) + shifts = tl.reshape(shifts, (BLOCK_SIZE_K, BLOCK_SIZE_N)) + + # Offsets and masks. + offsets_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + masks_am = offsets_am < M + + offsets_bn = pid_n * (BLOCK_SIZE_N // 8) + tl.arange(0, BLOCK_SIZE_N // 8) + masks_bn = offsets_bn < N // 8 + + offsets_zn = pid_n * (BLOCK_SIZE_N // 8) + tl.arange(0, BLOCK_SIZE_N // 8) + masks_zn = offsets_zn < N // 8 + + offsets_sn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + masks_sn = offsets_sn < N + + offsets_k = pid_z * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + offsets_a = K * offsets_am[:, None] + offsets_k[None, :] + offsets_b = (N // 8) * offsets_k[:, None] + offsets_bn[None, :] + + a_ptrs = a_ptr + offsets_a + b_ptrs = b_ptr + offsets_b + + # NOTE: Use this in TRITON_INTERPRET=1 mode instead of tl.cdiv + # block_offset = BLOCK_SIZE_K * SPLIT_K + # for k in range(0, (K + block_offset - 1) // (block_offset)): + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)): + masks_k = offsets_k < K + masks_a = masks_am[:, None] & masks_k[None, :] + a = tl.load(a_ptrs, mask=masks_a, other=0.0) + + masks_b = masks_k[:, None] & masks_bn[None, :] + b = tl.load(b_ptrs, mask=masks_b, other=0.0) + b = tl.interleave(b, b) + b = tl.interleave(b, b) + b = tl.interleave(b, b) + + # Dequantize b. + offsets_szk = ( + BLOCK_SIZE_K * SPLIT_K * k + pid_z * BLOCK_SIZE_K + ) // group_size + tl.arange(0, 1) + offsets_z = (N // 8) * offsets_szk[:, None] + offsets_zn[None, :] + masks_zk = offsets_szk < K // group_size + masks_z = masks_zk[:, None] & masks_zn[None, :] + zeros_ptrs = zeros_ptr + offsets_z + zeros = tl.load(zeros_ptrs, mask=masks_z, other=0.0) + zeros = tl.interleave(zeros, zeros) + zeros = tl.interleave(zeros, zeros) + zeros = tl.interleave(zeros, zeros) + zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N)) + + offsets_s = N * offsets_szk[:, None] + offsets_sn[None, :] + masks_sk = offsets_szk < K // group_size + masks_s = masks_sk[:, None] & masks_sn[None, :] + scales_ptrs = scales_ptr + offsets_s + scales = tl.load(scales_ptrs, mask=masks_s, other=0.0) + scales = tl.broadcast_to(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N)) + + b = (b >> shifts) & 0xF + zeros = (zeros >> shifts) & 0xF + b = (b - zeros) * scales + b = b.to(c_ptr.type.element_ty) + + # Accumulate results. + accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype) + + offsets_k += BLOCK_SIZE_K * SPLIT_K + a_ptrs += BLOCK_SIZE_K * SPLIT_K + b_ptrs += BLOCK_SIZE_K * SPLIT_K * (N // 8) + + c = accumulator.to(c_ptr.type.element_ty) + 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_ptr + pid_z * N * M + N * offs_cm[:, None] + offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +# qweights - [K , M // 8], int32 +# scales - [K // G, M ], float16 +# zeros - [K // G, M // 8], int32 +def awq_dequantize_triton( + qweight: torch.Tensor, + scales: torch.Tensor, + zeros: torch.Tensor, + block_size_x: int = 32, + block_size_y: int = 32, +) -> torch.Tensor: + K = qweight.shape[0] + M = scales.shape[1] + group_size = qweight.shape[0] // scales.shape[0] + + assert K > 0 and M > 0 + assert scales.shape[0] == K // group_size and scales.shape[1] == M + assert zeros.shape[0] == K // group_size and zeros.shape[1] == M // 8 + assert group_size <= K + assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K + + # Result tensor: + # number of rows = same as input tensor + # number of cols = 8 x input tensor num cols + result = torch.empty( + qweight.shape[0], + qweight.shape[1] * 8, + device=qweight.device, + dtype=scales.dtype, + ) + + Y = qweight.shape[0] # num rows + X = qweight.shape[1] # num cols + + grid = lambda META: ( + triton.cdiv(X, META["BLOCK_SIZE_X"]), + triton.cdiv(Y, META["BLOCK_SIZE_Y"]), + ) + awq_dequantize_kernel[grid]( + qweight, + scales, + zeros, + group_size, + result, + X, + Y, + BLOCK_SIZE_X=block_size_x, + BLOCK_SIZE_Y=block_size_y, + ) + + return result + + +# input - [M, K] +# qweight - [K, N // 8] +# qzeros - [K // G, N // 8] +# scales - [K // G, N] +# split_k_iters - parallelism along K-dimension, int, power of 2. +def awq_gemm_triton( + input: torch.Tensor, + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor, + split_k_iters: int, + block_size_m: int = 32, + block_size_n: int = 32, + block_size_k: int = 32, +) -> torch.Tensor: + M, K = input.shape + N = qweight.shape[1] * 8 + group_size = qweight.shape[0] // qzeros.shape[0] + + assert N > 0 and K > 0 and M > 0 + assert qweight.shape[0] == K and qweight.shape[1] == N // 8 + assert qzeros.shape[0] == K // group_size and qzeros.shape[1] == N // 8 + assert scales.shape[0] == K // group_size and scales.shape[1] == N + assert split_k_iters & (split_k_iters - 1) == 0 and split_k_iters != 0 + assert split_k_iters <= 32 + assert group_size <= K + assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K + + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + split_k_iters, + ) + + result = torch.zeros((split_k_iters, M, N), dtype=scales.dtype, device=input.device) + + # A = input, B = qweight, C = result + # A = M x K, B = K x N, C = M x N + awq_gemm_kernel[grid]( + input, + qweight, + result, + qzeros, + scales, + M, + N, + K, + group_size, + BLOCK_SIZE_M=block_size_m, + BLOCK_SIZE_N=block_size_n, + BLOCK_SIZE_K=block_size_k, + SPLIT_K=split_k_iters, + ) + + result = result.sum(0) + + return result diff --git a/python/sglang/srt/layers/quantization/base_config.py b/python/sglang/srt/layers/quantization/base_config.py new file mode 100644 index 000000000..4a5b7905e --- /dev/null +++ b/python/sglang/srt/layers/quantization/base_config.py @@ -0,0 +1,216 @@ +# Adapted from https://raw.githubusercontent.com/vllm-project/vllm/v0.5.5/vllm/model_executor/layers/quantization/base_config.py +from __future__ import annotations + +import inspect +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type + +import torch +from torch import nn + +if TYPE_CHECKING: + from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig + from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput + + +class QuantizeMethodBase(ABC): + """Base class for different quantized methods.""" + + @abstractmethod + def create_weights( + self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs + ): + """Create weights for a layer. + + The weights will be set as attributes of the layer.""" + raise NotImplementedError() + + @abstractmethod + def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: + """Apply the weights in layer to the input tensor. + + Expects create_weights to have been called before on the layer.""" + raise NotImplementedError() + + def process_weights_after_loading(self, layer: nn.Module) -> None: + """Process the weight after loading. + + This can be used for example, to transpose weights for computation. + """ + return + + +class LinearMethodBase(QuantizeMethodBase): + """Base class for different (maybe quantized) linear methods.""" + + @abstractmethod + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + """Create weights for a linear layer. + The weights will be set as attributes of the layer. + + Args: + layer: The layer that is using the LinearMethodBase factory. + input_size_per_partition: Size of the weight input dim on rank X. + output_partition_sizes: Sizes of the output dim of each logical + weight on rank X. E.g., output_partition_sizes for QKVLinear + is a list contains the width of Wq, Wk, Wv on rank X. + input_size: Size of the input dim of the weight across all ranks. + output_size: Size of the output dim of the weight across all ranks. + params_dtype: Datatype of the parameters. + """ + raise NotImplementedError() + + @abstractmethod + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Apply the weights in layer to the input tensor. + Expects create_weights to have been called before on the layer.""" + raise NotImplementedError() + + +class FusedMoEMethodBase(QuantizeMethodBase): + + @abstractmethod + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + raise NotImplementedError + + @abstractmethod + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + raise NotImplementedError + + @abstractmethod + def apply( + self, + layer: torch.nn.Module, + dispatch_output: DispatchOutput, + ) -> CombineInput: + raise NotImplementedError + + +class QuantizationConfig(ABC): + """Base class for quantization configs.""" + + def __init__(self): + super().__init__() + # mapping is updated by models as they initialize + self.packed_modules_mapping: Dict[str, List[str]] = dict() + + @abstractmethod + def get_name(self) -> str: + """Name of the quantization method.""" + raise NotImplementedError() + + @abstractmethod + def get_supported_act_dtypes(self) -> List[torch.dtype]: + """List of supported activation dtypes.""" + raise NotImplementedError() + + @classmethod + @abstractmethod + def get_min_capability(cls) -> int: + """Minimum GPU capability to support the quantization method. + + E.g., 70 for Volta, 75 for Turing, 80 for Ampere. + This requirement is due to the custom CUDA kernels used by the + quantization method. + """ + raise NotImplementedError() + + @staticmethod + @abstractmethod + def get_config_filenames() -> List[str]: + """List of filenames to search for in the model directory.""" + raise NotImplementedError() + + @classmethod + @abstractmethod + def from_config(cls, config: Dict[str, Any]) -> "QuantizationConfig": + """Create a config class from the model's quantization config.""" + raise NotImplementedError() + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + """ + Detects if this quantization method can support a given checkpoint + format by overriding the user specified quantization method -- + this method should only be overwritten by subclasses in exceptional + circumstances + """ + return None + + @staticmethod + def get_from_keys(config: Dict[str, Any], keys: List[str]) -> Any: + """Get a value from the model's quantization config.""" + for key in keys: + if key in config: + return config[key] + raise ValueError( + f"Cannot find any of {keys} in the model's " "quantization config." + ) + + @staticmethod + def get_from_keys_or(config: Dict[str, Any], keys: List[str], default: Any) -> Any: + """Get a optional value from the model's quantization config.""" + try: + return QuantizationConfig.get_from_keys(config, keys) + except ValueError: + return default + + @abstractmethod + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + """Get the quantize method to use for the quantized layer. + + Args: + layer: The layer for the quant method. + prefix: The full name of the layer in the state dict + Returns: + The quantize method. None if the given layer doesn't support quant + method. + """ + raise NotImplementedError() + + @abstractmethod + def get_scaled_act_names(self) -> List[str]: + """Returns the activation function names that should be post-scaled. + + For now, this is only used by AWQ. + """ + raise NotImplementedError() + + +def method_has_implemented_embedding(method_class: Type[QuantizeMethodBase]) -> bool: + """ + Not all quant methods have embedding implemented, so we need to check that + it exists for our given method. We check this by making sure the function + has been changed from the base implementation. + """ + base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding", None) + class_embedding = inspect.getattr_static(method_class, "embedding", None) + + return class_embedding is not None and class_embedding is not base_embedding diff --git a/python/sglang/srt/layers/quantization/blockwise_int8.py b/python/sglang/srt/layers/quantization/blockwise_int8.py new file mode 100644 index 000000000..60d4e3929 --- /dev/null +++ b/python/sglang/srt/layers/quantization/blockwise_int8.py @@ -0,0 +1,380 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/layers/quantization/fp8.py + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import torch +from torch.nn import Module + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.parameter import BlockQuantScaleParameter, ModelWeightParameter +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.int8_utils import apply_w8a8_block_int8_linear +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.layers.quantization.utils import is_layer_skipped +from sglang.srt.utils import set_weight_attrs + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +ACTIVATION_SCHEMES = ["static", "dynamic"] + +logger = logging.getLogger(__name__) + + +class BlockInt8Config(QuantizationConfig): + """Config class for INT8.""" + + def __init__( + self, + is_checkpoint_int8_serialized: bool = False, + activation_scheme: str = "dynamic", + ignored_layers: Optional[List[str]] = None, + weight_block_size: List[int] = None, + ) -> None: + self.is_checkpoint_int8_serialized = is_checkpoint_int8_serialized + if is_checkpoint_int8_serialized: + logger.warning( + "Detected int8 checkpoint. Please note that the " + "format is experimental and subject to change." + ) + if activation_scheme not in ACTIVATION_SCHEMES: + raise ValueError(f"Unsupported activation scheme {activation_scheme}") + self.activation_scheme = activation_scheme + self.ignored_layers = ignored_layers or [] + if weight_block_size is not None: + if not is_checkpoint_int8_serialized: + raise ValueError( + f"The block-wise quantization only supports int8-serialized checkpoint for now." + ) + if len(weight_block_size) != 2: + raise ValueError( + f"The quantization block size of weight must have 2 dimensions, but got {len(weight_block_size)} dimensions." + ) + if activation_scheme != "dynamic": + raise ValueError( + f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme." + ) + self.weight_block_size = weight_block_size + + @classmethod + def get_name(cls) -> str: + return "blockwise_int8" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> BlockInt8Config: + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_int8_serialized = "int8" in quant_method + activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) + return cls( + is_checkpoint_int8_serialized=is_checkpoint_int8_serialized, + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + weight_block_size=weight_block_size, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + if is_layer_skipped(prefix, self.ignored_layers): + return UnquantizedLinearMethod() + return BlockInt8LinearMethod(self) + elif isinstance(layer, FusedMoE): + return BlockInt8MoEMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class BlockInt8LinearMethod(LinearMethodBase): + """Linear method for INT8. + Supports loading INT8 checkpoints with static weight scale and + dynamic activation scale. + + Limitations: + Only support block-wise int8 quantization and int8 checkpoint + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: BlockInt8Config): + self.quant_config = quant_config + assert self.quant_config.weight_block_size is not None + assert self.quant_config.is_checkpoint_int8_serialized + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + tp_size = get_tensor_model_parallel_world_size() + + block_n, block_k = ( + self.quant_config.weight_block_size[0], + self.quant_config.weight_block_size[1], + ) + # Required by row parallel + if tp_size > 1 and input_size // input_size_per_partition == tp_size: + if input_size_per_partition % block_k != 0: + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + # Required by column parallel or enabling merged weights + if (tp_size > 1 and output_size // output_size_per_partition == tp_size) or len( + output_partition_sizes + ) > 1: + for output_partition_size in output_partition_sizes: + if output_partition_size % block_n != 0: + raise ValueError( + f"Weight output_partition_size = " + f"{output_partition_size} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + + layer.logical_widths = output_partition_sizes + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # WEIGHT + weight_dtype = ( + torch.int8 + if self.quant_config.is_checkpoint_int8_serialized + else params_dtype + ) + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, input_size_per_partition, dtype=weight_dtype + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + + scale = BlockQuantScaleParameter( + data=torch.empty( + (output_size_per_partition + block_n - 1) // block_n, + (input_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale_inv", scale) + + # INPUT ACTIVATION SCALE + assert self.quant_config.activation_scheme == "dynamic" + layer.register_parameter("input_scale", None) + + def process_weights_after_loading(self, layer: Module) -> None: + # Block quant doesn't need to process weights after loading + # Use torch Parameter to avoid cuda graph capturing issue + layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False) + layer.weight_scale_inv = torch.nn.Parameter( + layer.weight_scale_inv.data, requires_grad=False + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_w8a8_block_int8_linear( + input=x, + weight=layer.weight, + block_size=self.quant_config.weight_block_size, + weight_scale=layer.weight_scale_inv, + input_scale=None, + bias=bias, + ) + + +class BlockInt8MoEMethod(FusedMoEMethodBase): + """MoE method for INT8. + Supports loading INT8 checkpoints with static weight scale and + dynamic activation scale. + + Limitations: + Only support block-wise int8 quantization and int8 checkpoint + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: BlockInt8Config): + self.quant_config = quant_config + assert self.quant_config.weight_block_size is not None + assert self.quant_config.is_checkpoint_int8_serialized + + def create_weights( + self, + layer: Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + if self.quant_config.is_checkpoint_int8_serialized: + params_dtype = torch.int8 + tp_size = get_tensor_model_parallel_world_size() + + block_n, block_k = ( + self.quant_config.weight_block_size[0], + self.quant_config.weight_block_size[1], + ) + # NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n. + # Required by column parallel or enabling merged weights + if intermediate_size_per_partition % block_n != 0: + raise ValueError( + f"The output_size of gate's and up's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + if tp_size > 1: + # Required by row parallel + if intermediate_size_per_partition % block_k != 0: + raise ValueError( + f"The input_size of down's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * ((intermediate_size_per_partition + block_n - 1) // block_n), + (hidden_size + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + (hidden_size + block_n - 1) // block_n, + (intermediate_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale_inv", w13_weight_scale) + layer.register_parameter("w2_weight_scale_inv", w2_weight_scale) + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # INPUT_SCALES + assert self.quant_config.activation_scheme == "dynamic" + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer: Module) -> None: + # Block quant doesn't need to process weights after loading + return + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + use_int8_w8a8=True, + w13_scale=layer.w13_weight_scale_inv, + w2_scale=layer.w2_weight_scale_inv, + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + block_shape=self.quant_config.weight_block_size, + ) + + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/README.md b/python/sglang/srt/layers/quantization/compressed_tensors/README.md new file mode 100644 index 000000000..e8caf538b --- /dev/null +++ b/python/sglang/srt/layers/quantization/compressed_tensors/README.md @@ -0,0 +1,6 @@ +# quantization compressed_tensors module + +To support compressed_tensors format quantization models, we adapted https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors into SGLang. + + +For practical purposes, we have only applied the compressed_tensors format of `w8a8_fp8`. If you have requirements for other formats, you can submit an issue through this [link](https://github.com/sgl-project/sglang/issues). diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/__init__.py b/python/sglang/srt/layers/quantization/compressed_tensors/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py new file mode 100644 index 000000000..8afc15a73 --- /dev/null +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -0,0 +1,666 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/v0.8.2/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import logging +from contextlib import suppress +from typing import Any, Dict, List, Literal, NamedTuple, Optional, Tuple, cast + +import torch +from compressed_tensors.config import ( + CompressionFormat, + SparsityCompressionConfig, + SparsityStructure, +) +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) +from pydantic import BaseModel + +from sglang.srt.layers.quantization.base_config import ( + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + CompressedTensorsMoEMethod, +) +from sglang.srt.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsScheme, + CompressedTensorsW8A8Fp8, + CompressedTensorsW8A16Fp8, +) +from sglang.srt.layers.quantization.compressed_tensors.utils import ( + find_matched_target, + is_activation_quantization_format, + should_ignore_layer, +) +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod + +try: + from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( + WNA16_SUPPORTED_BITS, + CompressedTensorsWNA16, + ) + + VLLM_AVAILABLE = True +except ImportError: + VLLM_AVAILABLE = False + +logger = logging.getLogger(__name__) + +__all__ = ["CompressedTensorsLinearMethod"] + +SPARSITY_CONFIG_NAME: Literal["sparsity_config"] = "sparsity_config" +QUANTIZATION_SCHEME_MAP_TYPE = Dict[str, Optional[Dict[str, QuantizationArgs]]] + + +class DeviceCapability(NamedTuple): + major: int + minor: int + + def as_version_str(self) -> str: + return f"{self.major}.{self.minor}" + + def to_int(self) -> int: + """ + Express device capability as an integer ````. + + It is assumed that the minor version is always a single digit. + """ + assert 0 <= self.minor < 10 + return self.major * 10 + self.minor + + +class CompressedTensorsConfig(QuantizationConfig): + + def __init__( + self, + target_scheme_map: Dict[str, Any], + ignore: List[str], + quant_format: str, + sparsity_scheme_map: Dict[str, SparsityCompressionConfig], + sparsity_ignore_list: List[str], + kv_cache_scheme: Optional[Dict[str, Any]] = None, + config: Optional[Dict[str, Any]] = None, + packed_modules_mapping: Dict[str, List[str]] = {}, + ): + super().__init__() + self.ignore = ignore + self.quant_format = quant_format + # Map from [target -> scheme] + self.target_scheme_map = target_scheme_map + self.kv_cache_scheme = kv_cache_scheme + self.sparsity_scheme_map = sparsity_scheme_map + self.sparsity_ignore_list = sparsity_ignore_list + self.config = config + self.packed_modules_mapping = packed_modules_mapping + + def get_linear_method(self) -> CompressedTensorsLinearMethod: + return CompressedTensorsLinearMethod(self) + + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.float16, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 70 + + def get_name(self) -> str: + return "compressed_tensors" + + def get_scaled_act_names(self) -> List[str]: + return [] + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + + # Check if the layer is skipped for quantization. + # TODO (@robertgshaw2): support module names + if should_ignore_layer( + prefix, ignore=self.ignore, fused_mapping=self.packed_modules_mapping + ): + return UnquantizedLinearMethod() + if isinstance(layer, LinearBase): + scheme = self.get_scheme(layer=layer, layer_name=prefix) + if scheme is None: + return UnquantizedLinearMethod() + layer.scheme = scheme + return CompressedTensorsLinearMethod(self) + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, FusedMoE): + return CompressedTensorsMoEMethod.get_moe_method(self) + return None + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> CompressedTensorsConfig: + ignore: List[str] = cast(List[str], config.get("ignore", [])) + quant_format = cast(str, config.get("format")) + target_scheme_map = cls._quantization_scheme_map_from_config(config=config) + sparsity_scheme_map, sparsity_ignore_list = cls._parse_sparsity_config( + config=config + ) + packed_modules_mapping = config.get("packed_modules_mapping", {}) + + return cls( + target_scheme_map=target_scheme_map, + ignore=ignore, + quant_format=quant_format, + sparsity_scheme_map=sparsity_scheme_map, + sparsity_ignore_list=sparsity_ignore_list, + config=config, + packed_modules_mapping=packed_modules_mapping, + ) + + @classmethod + def _parse_sparsity_config( + cls, config: Dict[str, Any] + ) -> Tuple[Dict[str, SparsityCompressionConfig], List[str]]: + """ + :param config: The `quantization_config` dictionary from config.json + :return: A tuple with two elements + 1. A dictionary mapping target layer names to their corresponding + sparsity_config + 2. A list of layer names to ignore for sparsity + """ + if not (sparsity_config := config.get(SPARSITY_CONFIG_NAME)): + return dict(), [] + + sparsity_config = SparsityCompressionConfig.model_validate(sparsity_config) + sparse_scheme_map: Dict[str, SparsityCompressionConfig] = { + target: sparsity_config for target in sparsity_config.targets or list() + } + sparsity_ignore_list = sparsity_config.ignore or list() + return sparse_scheme_map, sparsity_ignore_list + + @classmethod + def _quantization_scheme_map_from_config( + cls, config: Dict[str, Any] + ) -> QUANTIZATION_SCHEME_MAP_TYPE: + """ + :param config: The `quantization_config` dictionary from config.json + :return: A dictionary mapping target layer names to their corresponding + quantization_args for weights and input activations + """ + target_scheme_map: Dict[str, Any] = dict() + quant_format = cast(str, config.get("format")) + + # The quant_config has multiple config_groups, each containing + # an input_activations key with details about how the activations are + # quantized, a weights key indicating how the weights are quantized, + # and a list of targets under the `targets` key, dictating which + # layers are impacted by the quantization details. The quantization + # details follow the structure defined by the QuantizationArgs + # pydantic model, which is used to verify the structure of the + # quant_config and also store the details for later use. + + config_groups = config.get("config_groups", dict()) + for _, quant_config in config_groups.items(): + targets = quant_config.get("targets") + for target in targets: + target_scheme_map[target] = {} + target_scheme_map[target]["weights"] = QuantizationArgs.model_validate( + quant_config.get("weights") + ) + + target_scheme_map[target]["input_activations"] = None + if is_activation_quantization_format(quant_format): + input_activations = quant_config.get("input_activations") + # The only case where we have activation quant supported + # but no input_activations provided in the config + # should be w8a16fp8 w8a16fp8 can also run for cases where + # there is an input_quant but it is ignored + if not input_activations: + assert ( + target_scheme_map[target]["weights"].type + == QuantizationType.FLOAT + ) + else: + target_scheme_map[target]["input_activations"] = ( + QuantizationArgs.model_validate( # noqa: E501 + quant_config.get("input_activations") + ) + ) + return target_scheme_map + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + def _check_scheme_supported(self, min_capability: int, error: bool = True) -> bool: + capability_tuple = DeviceCapability(*torch.cuda.get_device_capability()) + + if capability_tuple is not None: + capability = capability_tuple.to_int() + supported = capability >= min_capability + if error and not supported: + raise RuntimeError( + "Quantization scheme is not supported for ", + f"the current GPU. Min capability: {min_capability}. ", + f"Current capability: {capability}.", + ) + return supported + else: + return False + + def _is_static_tensor_w8a8( + self, weight_quant: BaseModel, input_quant: BaseModel + ) -> bool: + is_8_bits = weight_quant.num_bits == input_quant.num_bits == 8 + weight_strategy = ( + weight_quant.strategy == QuantizationStrategy.TENSOR.value + or weight_quant.strategy == QuantizationStrategy.CHANNEL.value + ) + is_tensor = ( + weight_strategy + and input_quant.strategy == QuantizationStrategy.TENSOR.value + ) + is_static = not weight_quant.dynamic and not input_quant.dynamic + + # Both symmetric and asymmetric input quantization supported. + # Only symmetric weight quantization supported. + return is_8_bits and is_tensor and weight_quant.symmetric and is_static + + def _is_dynamic_token_w8a8( + self, weight_quant: BaseModel, input_quant: BaseModel + ) -> bool: + is_8_bits = weight_quant.num_bits == input_quant.num_bits == 8 + weight_strategy = ( + weight_quant.strategy == QuantizationStrategy.TENSOR.value + or weight_quant.strategy == QuantizationStrategy.CHANNEL.value + ) + is_token = ( + weight_strategy and input_quant.strategy == QuantizationStrategy.TOKEN.value + ) + is_dynamic = not weight_quant.dynamic and input_quant.dynamic + + # Both symmetric and asymmetric input quantization supported. + # Only symmetric weight quantization supported. + return is_8_bits and is_token and weight_quant.symmetric and is_dynamic + + def _is_fp8_w8a8(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool: + # Confirm weights and activations quantized. + if weight_quant is None or input_quant is None: + return False + + # Confirm weight scheme is supported. + is_floating_point = ( + weight_quant.type == QuantizationType.FLOAT + and input_quant.type == QuantizationType.FLOAT + ) + is_symmetric_weight = weight_quant.symmetric + is_static_weight = not weight_quant.dynamic + is_per_tensor_or_channel_weight = weight_quant.strategy in [ + QuantizationStrategy.TENSOR, + QuantizationStrategy.CHANNEL, + ] + if not ( + is_floating_point + and is_symmetric_weight + and is_static_weight + and is_per_tensor_or_channel_weight + ): + return False + + # Dynamic quantization is always supported if weights supported. + if input_quant.dynamic: + return True + + # Confirm activation scheme is supported. + is_symmetric_activation = input_quant.symmetric + is_per_tensor_activation = input_quant.strategy == QuantizationStrategy.TENSOR + return is_symmetric_activation and is_per_tensor_activation + + def _is_fp8_w8a16(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool: + # Confirm weights quantized. + if weight_quant is None: + return False + + # Confirm we have floating points. + if weight_quant.type != QuantizationType.FLOAT: + return False + + # Confirm weight scheme is supported. + is_symmetric_weight = weight_quant.symmetric + is_static_weight = not weight_quant.dynamic + is_per_tensor_or_channel_weight = weight_quant.strategy in [ + QuantizationStrategy.TENSOR, + QuantizationStrategy.CHANNEL, + ] + if not ( + is_symmetric_weight + and is_static_weight # noqa: SIM103 + and is_per_tensor_or_channel_weight + ): + return False + + # All conditions satisfied. + return True + + def _is_wNa16_group_channel( + self, weight_quant: BaseModel, input_quant: BaseModel + ) -> bool: + input_quant_none = input_quant is None + is_symmetric = weight_quant.symmetric + is_channel_group = ( + weight_quant.strategy == QuantizationStrategy.CHANNEL.value + or weight_quant.strategy == QuantizationStrategy.GROUP.value + ) + is_static = not weight_quant.dynamic + + return is_channel_group and input_quant_none and is_symmetric and is_static + + def _get_scheme_from_parts( + self, weight_quant: BaseModel, input_quant: BaseModel + ) -> CompressedTensorsScheme: + + # Detect If Mixed Precision + if self._is_wNa16_group_channel(weight_quant, input_quant): + if not VLLM_AVAILABLE: + raise ImportError( + "vllm is not installed, to use CompressedTensorsW4A16Sparse24 and CompressedTensorsWNA16, please install vllm" + ) + if ( + self.quant_format == CompressionFormat.marlin_24.value + and weight_quant.num_bits in W4A16SPARSE24_SUPPORTED_BITS + ): + return CompressedTensorsW4A16Sparse24( + strategy=weight_quant.strategy, + num_bits=weight_quant.num_bits, + group_size=weight_quant.group_size, + ) + if ( + self.quant_format == CompressionFormat.pack_quantized.value + and weight_quant.num_bits in WNA16_SUPPORTED_BITS + ): + return CompressedTensorsWNA16( + num_bits=weight_quant.num_bits, + strategy=weight_quant.strategy, + group_size=weight_quant.group_size, + actorder=weight_quant.actorder, + ) + + if is_activation_quantization_format(self.quant_format): + if self._is_fp8_w8a8(weight_quant, input_quant): + is_fp8_w8a8_supported = self._check_scheme_supported( + CompressedTensorsW8A8Fp8.get_min_capability(), error=False + ) + if is_fp8_w8a8_supported: + return CompressedTensorsW8A8Fp8( + strategy=weight_quant.strategy, + is_static_input_scheme=( + input_quant and not input_quant.dynamic + ), + ) + else: + # note: input_quant will be present for converted models; + # will be ignored during inference post loading + return CompressedTensorsW8A16Fp8( + strategy=weight_quant.strategy, + is_static_input_scheme=not input_quant.dynamic, + ) + + # note: input_quant can be None + if self._is_fp8_w8a16(weight_quant, input_quant): + if not VLLM_AVAILABLE: + raise ImportError( + "vllm is not installed, to use CompressedTensorsW8A16Fp8, please install vllm" + ) + is_static_input_scheme = input_quant and not input_quant.dynamic + return CompressedTensorsW8A16Fp8( + strategy=weight_quant.strategy, + is_static_input_scheme=is_static_input_scheme, + ) + + if self._is_static_tensor_w8a8(weight_quant, input_quant): + return CompressedTensorsW8A8Int8( + strategy=weight_quant.strategy, + is_static_input_scheme=True, + input_symmetric=input_quant.symmetric, + ) + + if self._is_dynamic_token_w8a8(weight_quant, input_quant): + return CompressedTensorsW8A8Int8( + strategy=weight_quant.strategy, + is_static_input_scheme=False, + input_symmetric=input_quant.symmetric, + ) + + raise NotImplementedError("No compressed-tensors compatible scheme was found.") + + def get_scheme( + self, layer: torch.nn.Module, layer_name: Optional[str] = None + ) -> Optional[CompressedTensorsScheme]: + """ + compressed-tensors supports non uniform in the following way: + + targets of config_groups: There can be N config_groups which each + have a quantization scheme. Each config_group has a list of targets + which can be a full layer_name, a regex for a layer_name, or + an nn.Module name. + + Detect whether a layer_name is found in any target and + use the quantization scheme corresponding to the matched target + to select the CompressedTensorsScheme used for infernece. + """ + + # Find the "target" in the compressed-tensors config + # that our layer conforms to. + # TODO (@robertgshaw): add compressed-tensors as dep + # so we do not have to re-write these functions + # need to make accelerate optional in ct to do this + + # Will be empty for models with only sparsity + weight_quant = input_quant = None + if self.target_scheme_map: + matched_target = find_matched_target( + layer_name=layer_name, + module=layer, + targets=self.target_scheme_map.keys(), + fused_mapping=self.packed_modules_mapping, + ) + + scheme_dict = self.target_scheme_map[matched_target] + weight_quant = scheme_dict.get("weights") + input_quant = scheme_dict.get("input_activations") + + # Find the sparsity scheme of the layer + # assume that fused layers inerhit first component's sparsity scheme + sparsity_targets = self.sparsity_scheme_map.keys() - set( + self.sparsity_ignore_list + ) + sparsity_scheme: Optional[SparsityCompressionConfig] = None + with suppress(ValueError): + matched_target = find_matched_target( + layer_name=layer_name, + module=layer, + targets=sparsity_targets, + fused_mapping=self.packed_modules_mapping, + ) + sparsity_scheme = self.sparsity_scheme_map[matched_target] + + if self.supports_cutlass_24( + weight_quant=weight_quant, + input_quant=input_quant, + sparsity_scheme=sparsity_scheme, + ): + if not VLLM_AVAILABLE: + raise ImportError( + "vllm is not installed, to use CompressedTensors24, please install vllm" + ) + # Have a valid sparsity scheme + # Validate layer is supported by Cutlass 2:4 Kernel + model_compression_config = ( + None + if sparsity_scheme is None or sparsity_scheme.format == "dense" + else self.config + ) + + scheme = CompressedTensors24( + quantized=weight_quant is not None or input_quant is not None, + weight_quant=weight_quant, + input_quant=input_quant, + model_compression_config=model_compression_config, + ) + elif weight_quant is None: + logger.warning_once( + "Acceleration for non-quantized schemes is " + "not supported by Compressed Tensors. " + "Falling back to UnquantizedLinearMethod" + ) + return None + + else: + # Find the quant_scheme + scheme = self._get_scheme_from_parts( # type: ignore + weight_quant=weight_quant, + input_quant=input_quant, + ) + + # Raise error if device does not support the scheme + # (e.g. fp8 needs ada lovelace) + self._check_scheme_supported(scheme.get_min_capability()) + logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name) + return scheme + + def get_cache_scale(self, name: str) -> Optional[str]: + """ + Check whether the param name matches the format for k/v cache scales + in compressed-tensors. If this is the case, return its equivalent + param name expected by vLLM + + :param name: param name + :return: matching param name for KV cache scale in vLLM + """ + if name.endswith(".output_scale") and ".k_proj" in name: + return name.replace(".k_proj.output_scale", ".attn.k_scale") + if name.endswith(".output_scale") and ".v_proj" in name: + return name.replace(".v_proj.output_scale", ".attn.v_scale") + # If no matches, return None + return None + + @staticmethod + def supports_cutlass_24( + weight_quant: Optional[QuantizationArgs], + input_quant: Optional[QuantizationArgs], + sparsity_scheme: Optional[SparsityCompressionConfig] = None, + ) -> bool: + """ + Check if the layer is supported by the Cutlass 2:4 Kernel + Conditions: + - Overarching condition: Sparsity Structure is 2:4 + - Unquantized cases are supported + - Weight only quantization is not-supported + - Supported weight quantization strategies are TENSOR and CHANNEL + - Supported input quantization strategies are TENSOR and TOKEN + - Only 8 bit quantization is supported + + :return: True if the layer is supported by the Cutlass 2:4 Kernel + False otherwise + """ + if sparsity_scheme is None: + return False + + is_valid_sparsity_structure: bool = ( + sparsity_scheme.sparsity_structure == SparsityStructure.TWO_FOUR.value + ) + + valid_compressors = { + CompressionFormat.dense.value, + CompressionFormat.sparse_24_bitmask.value, + } + + is_valid_sparsity = ( + is_valid_sparsity_structure and sparsity_scheme.format in valid_compressors + ) + + if not is_valid_sparsity: + return False + + # Unquantized cases are supported + if weight_quant is None and input_quant is None: + return True + + # Weight only quantization is not-supported + if weight_quant is not None and input_quant is None: + return False + + supported_weight_quant_strategies = [ + QuantizationStrategy.TENSOR.value, + QuantizationStrategy.CHANNEL.value, + ] + + assert weight_quant is not None + assert input_quant is not None + if weight_quant.strategy not in supported_weight_quant_strategies: + return False + + supported_input_quant_strategies = [ + QuantizationStrategy.TENSOR.value, + QuantizationStrategy.TOKEN.value, + ] + + if input_quant.strategy not in supported_input_quant_strategies: + return False + + return weight_quant.num_bits == input_quant.num_bits == 8 + + +class CompressedTensorsLinearMethod(LinearMethodBase): + + def __init__(self, quantization_config: CompressedTensorsConfig): + self.quantization_config = quantization_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + """ + Use the CompressedTensorsScheme associated with each layer to create + the necessary parameters for the layer. See LinearMethodBase for param + details + """ + weight_loader = extra_weight_attrs.get("weight_loader") + layer.scheme.create_weights( + layer=layer, + input_size=input_size, + input_size_per_partition=input_size_per_partition, + output_partition_sizes=output_partition_sizes, + output_size=output_size, + params_dtype=params_dtype, + weight_loader=weight_loader, + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + """ + Use the output of create_weights and the CompressedTensorsScheme + associated with the layer to apply the forward pass with the + layer input. See LinearMethodBase for param details + + """ + + scheme = layer.scheme + if scheme is None: + raise ValueError("A scheme must be defined for each layer") + return scheme.apply_weights(layer, x, bias=bias) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py new file mode 100644 index 000000000..e2ff25e68 --- /dev/null +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -0,0 +1,692 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/v0.8.2/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import enum +import logging +from enum import Enum +from typing import TYPE_CHECKING, List, Optional + +import torch +from compressed_tensors import CompressionFormat +from compressed_tensors.quantization import QuantizationStrategy + +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz, scaled_fp8_quant +from sglang.srt.layers.quantization.fp8_utils import normalize_e4m3fn_to_e4m3fnuz +from sglang.srt.layers.quantization.utils import ( + all_close_1d, + per_tensor_dequantize, + replace_parameter, +) +from sglang.srt.utils import ( + get_bool_env_var, + is_cpu, + is_cuda, + is_hip, + is_npu, + set_weight_attrs, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, + ) + +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _use_aiter: + from aiter.ops.shuffle import shuffle_weight + + from sglang.srt.layers.moe.rocm_moe_utils import rocm_fused_experts_tkw1 + +try: + import vllm + + VLLM_AVAILABLE = True +except ImportError: + VLLM_AVAILABLE = False + +logger = logging.getLogger(__name__) + + +class GPTQMarlinState(Enum): + REPACK = enum.auto() + READY = enum.auto() + + +__all__ = [ + "CompressedTensorsMoEMethod", + "CompressedTensorsW8A8Fp8MoEMethod", + "CompressedTensorsWNA16MoEMethod", +] + + +class CompressedTensorsMoEMethod(FusedMoEMethodBase): + def __new__(cls, *args, **kwargs): + if cls is CompressedTensorsMoEMethod: + return super().__new__(cls) + return super().__new__(cls) + + @staticmethod + def get_moe_method( + quant_config: CompressedTensorsConfig, + ) -> "CompressedTensorsMoEMethod": + # TODO: @dsikka: refactor this to use schemes as other kernels + # are supported + check if the layer is being ignored. + weight_quant = quant_config.target_scheme_map["Linear"].get("weights") + input_quant = quant_config.target_scheme_map["Linear"].get("input_activations") + + if quant_config._is_wNa16_group_channel(weight_quant, input_quant): + if not VLLM_AVAILABLE: + raise ImportError( + "vllm is not installed, to use CompressedTensorsWNA16MoEMethod, please install vllm." + ) + return CompressedTensorsWNA16MoEMethod(quant_config) + elif quant_config._is_fp8_w8a8(weight_quant, input_quant): + return CompressedTensorsW8A8Fp8MoEMethod(quant_config) + else: + raise RuntimeError( + f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}" + ) + + +class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): + + def __init__(self, quant_config: CompressedTensorsConfig): + self.quant_config = quant_config + self.weight_quant = self.quant_config.target_scheme_map["Linear"].get("weights") + self.input_quant = self.quant_config.target_scheme_map["Linear"].get( + "input_activations" + ) + + self.static_input_scales = not self.input_quant.dynamic + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + params_dtype = torch.float8_e4m3fn + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + # per-tensor quantization + if self.weight_quant.strategy == QuantizationStrategy.TENSOR: + # Allocate 2 scales for w1 and w3 respectively. + # They will be combined to a single scale after weight loading. + w13_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + weight_quant_method = FusedMoeWeightScaleSupported.TENSOR.value + elif self.weight_quant.strategy == QuantizationStrategy.CHANNEL: + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * intermediate_size_per_partition, + 1, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + weight_quant_method = FusedMoeWeightScaleSupported.CHANNEL.value + else: + raise ValueError( + f"Unsupported weight quantization strategy: {self.weight_quant.strategy}" + ) + + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + # Add the quantization method used (per tensor/grouped/channel) + # to ensure the weight scales are loaded in properly + extra_weight_attrs.update({"quant_method": weight_quant_method}) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # INPUT_SCALES + if self.static_input_scales: + assert ( + self.input_quant.strategy == QuantizationStrategy.TENSOR + ), "Only per-tensor quantization is supported for static input scales" + w13_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + else: + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer: FusedMoE) -> None: + # Fp8 moe kernels require a single activation scale. + # We take the max of all the scales in case they differ. + if self.static_input_scales: + if layer.w13_input_scale is None or layer.w2_input_scale is None: + raise ValueError( + "QuantConfig has static quantization, but found " + "activation scales are None." + ) + if not all_close_1d(layer.w13_input_scale) or not all_close_1d( + layer.w2_input_scale + ): + logger.warning( + "Found input_scales that are not equal for " + "fp8 MoE layer. Using the maximum across experts " + "for each layer." + ) + layer.w13_input_scale = torch.nn.Parameter( + layer.w13_input_scale.max(), requires_grad=False + ) + layer.w2_input_scale = torch.nn.Parameter( + layer.w2_input_scale.max(), requires_grad=False + ) + + if is_fp8_fnuz(): + # Normalize the weights and scales + w13_weight, w13_weight_scale, w13_input_scale = ( + normalize_e4m3fn_to_e4m3fnuz( + layer.w13_weight, layer.w13_weight_scale, layer.w13_input_scale + ) + ) + w2_weight, w2_weight_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz( + layer.w2_weight, layer.w2_weight_scale, layer.w2_input_scale + ) + # Reset the parameter + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter( + w13_weight_scale, requires_grad=False + ) + if w13_input_scale is not None: + layer.w13_input_scale = torch.nn.Parameter( + w13_input_scale, requires_grad=False + ) + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter( + w2_weight_scale, requires_grad=False + ) + if w2_input_scale is not None: + layer.w2_input_scale = torch.nn.Parameter( + w2_input_scale, requires_grad=False + ) + if self.weight_quant.strategy == QuantizationStrategy.TENSOR: + # Fp8 moe kernel needs single weight scale for w13 per expert. + # We take the max then dequant and requant each expert. + assert layer.w13_weight_scale is not None + shard_size = layer.intermediate_size_per_partition + max_w13_scales = layer.w13_weight_scale.max(dim=1).values + for expert_id in range(layer.num_local_experts): + start = 0 + for shard_id in range(2): + dq_weight = per_tensor_dequantize( + layer.w13_weight[expert_id][start : start + shard_size, :], + layer.w13_weight_scale[expert_id][shard_id], + ) + ( + layer.w13_weight[expert_id][start : start + shard_size, :], + _, + ) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id]) + + start += shard_size + + layer.w13_weight_scale = torch.nn.Parameter( + max_w13_scales, requires_grad=False + ) + + if _use_aiter: + with torch.no_grad(): + # Pre-shuffle weights + layer.w13_weight = torch.nn.Parameter( + shuffle_weight(layer.w13_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + layer.w2_weight = torch.nn.Parameter( + shuffle_weight(layer.w2_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + moe_runner_config = self.moe_runner_config + + if ( + _use_aiter + and self.weight_quant.strategy == QuantizationStrategy.CHANNEL + and moe_runner_config.apply_router_weight_on_input + ): + topk_weights, topk_ids, _ = topk_output + output = rocm_fused_experts_tkw1( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=moe_runner_config.activation, + apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, + use_fp8_w8a8=True, + per_channel_quant=self.weight_quant.strategy + == QuantizationStrategy.CHANNEL, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + ) + return StandardCombineInput(hidden_states=output) + else: + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + use_fp8_w8a8=True, + per_channel_quant=self.weight_quant.strategy + == QuantizationStrategy.CHANNEL, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + ) + return self.runner.run(dispatch_output, quant_info) + + +class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): + + def __init__(self, quant_config: CompressedTensorsConfig): + self.quant_config = quant_config + # TODO: @dsikka: refactor this to use schemes as other kernels + # are supported + check if the layer is being ignored. + config = self.quant_config.target_scheme_map["Linear"].get("weights") + self.num_bits = config.num_bits + self.packed_factor = 32 // config.num_bits + self.strategy = config.strategy + self.group_size = config.group_size + self.actorder = config.actorder + assert config.symmetric, "Only symmetric quantization is supported for MoE" + + if not ( + self.quant_config.quant_format == CompressionFormat.pack_quantized.value + and self.num_bits in WNA16_SUPPORTED_BITS + ): + raise ValueError( + "For Fused MoE layers, only ", + f"{CompressionFormat.pack_quantized.value} ", + "is supported for the following bits: ", + f"{WNA16_SUPPORTED_BITS}", + ) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + + assert ( + params_dtype == torch.float16 + ), "float16 is required for MoE compressed models. Set dtype=torch.float16" # noqa: E501 + + # Will transpose the loaded weight along the + # intermediate and hidden dim sizes. Will + # shard for TP along the transposed dims + extra_weight_attrs.update( + {"is_transposed": True, "quant_method": self.strategy} + ) + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size // self.packed_factor, + 2 * intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition // self.packed_factor, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # In the case where we have actorder/g_idx, + # we do not partition the w2 scales + load_full_w2 = self.actorder and self.group_size != -1 + + if load_full_w2: + w2_scales_size = intermediate_size_per_partition * layer.moe_tp_size + else: + w2_scales_size = intermediate_size_per_partition + + self.is_k_full = (not self.actorder) or layer.moe_tp_size == 1 + + if self.strategy == "channel": + num_groups_w2 = num_groups_w13 = 1 + self.group_size = -1 + else: + num_groups_w2 = w2_scales_size // self.group_size + num_groups_w13 = hidden_size // self.group_size + + w13_scale = torch.nn.Parameter( + torch.ones( + num_experts, + num_groups_w13, + 2 * intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_scale) + set_weight_attrs(w13_scale, extra_weight_attrs) + + w2_scale = torch.nn.Parameter( + torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_scale) + set_weight_attrs(w2_scale, extra_weight_attrs) + set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2}) + + w2_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + layer.register_parameter("w2_weight_shape", w2_weight_shape) + set_weight_attrs(w2_weight_shape, extra_weight_attrs) + w13_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + + layer.register_parameter("w13_weight_shape", w13_weight_shape) + set_weight_attrs(w13_weight_shape, extra_weight_attrs) + + w13_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_g_idx", w13_g_idx) + set_weight_attrs(w13_g_idx, extra_weight_attrs) + + w2_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_g_idx", w2_g_idx) + set_weight_attrs(w2_g_idx, extra_weight_attrs) + + w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) + set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) + + w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) + set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) + + layer.a13_scale = None + layer.a2_scale = None + layer.marlin_state = GPTQMarlinState.REPACK + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + + def replace_tensor(name, new_t): + # It is important to use resize_() here since it ensures + # the same buffer is reused + getattr(layer, name).resize_(new_t.shape) + getattr(layer, name).copy_(new_t) + del new_t + + def get_scale_perms(num_bits: int): + scale_perm: List[int] = [] + for i in range(8): + scale_perm.extend([i + 8 * j for j in range(8)]) + scale_perm_single: List[int] = [] + for i in range(4): + scale_perm_single.extend( + [2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]] + ) + return scale_perm, scale_perm_single + + def marlin_permute_scales( + s: torch.Tensor, size_k: int, size_n: int, group_size: int, num_bits: int + ): + scale_perm, scale_perm_single = get_scale_perms(num_bits) + if group_size < size_k and group_size != -1: + s = s.reshape((-1, len(scale_perm)))[:, scale_perm] + else: + s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single] + s = s.reshape((-1, size_n)).contiguous() + return s + + def marlin_moe_permute_scales( + s: torch.Tensor, size_k: int, size_n: int, group_size: int, num_bits: int + ): + num_experts = s.shape[0] + output = torch.empty( + (num_experts, s.shape[1], s.shape[2]), device=s.device, dtype=s.dtype + ) + for e in range(num_experts): + output[e] = marlin_permute_scales( + s[e], size_k, size_n, group_size, num_bits + ) + return output + + size_k2 = layer.w2_weight_packed.shape[2] + size_k13 = layer.w13_weight_packed.shape[2] + + num_experts = layer.w13_weight_g_idx.shape[0] + device = layer.w13_weight_g_idx.device + + # when running models with grouped act order, + # resort to g_idx values provided in checkpoint + if self.actorder == "group": + w13_g_idx_sort_indices = torch.empty_like(layer.w13_weight_g_idx) + w2_g_idx_sort_indices = torch.empty_like(layer.w2_weight_g_idx) + w13_sorted_g_idx = torch.empty_like(layer.w13_weight_g_idx) + w2_sorted_g_idx = torch.empty_like(layer.w2_weight_g_idx) + + for e in range(num_experts): + w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_weight_g_idx[e]).to( + torch.int32 + ) + w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_weight_g_idx[e]).to( + torch.int32 + ) + w13_sorted_g_idx[e] = layer.w13_weight_g_idx[e][ + w13_g_idx_sort_indices[e] + ] + w2_sorted_g_idx[e] = layer.w2_weight_g_idx[e][w2_g_idx_sort_indices[e]] + + replace_parameter(layer, "w13_weight_g_idx", w13_sorted_g_idx) + replace_parameter(layer, "w2_weight_g_idx", w2_sorted_g_idx) + replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) + + else: + layer.w13_weight_g_idx = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_weight_g_idx = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + + from vllm import _custom_ops as vllm_ops + + marlin_w13_qweight = vllm_ops.gptq_marlin_moe_repack( + layer.w13_weight_packed, + layer.w13_g_idx_sort_indices, + layer.w13_weight_packed.shape[1] * self.packed_factor, + layer.w13_weight_packed.shape[2], + self.num_bits, + ) + replace_tensor("w13_weight_packed", marlin_w13_qweight) + marlin_w2_qweight = vllm_ops.gptq_marlin_moe_repack( + layer.w2_weight_packed, + layer.w2_g_idx_sort_indices, + layer.w2_weight_packed.shape[1] * self.packed_factor, + layer.w2_weight_packed.shape[2], + self.num_bits, + ) + replace_tensor("w2_weight_packed", marlin_w2_qweight) + # Repack scales + marlin_w13_scales = marlin_moe_permute_scales( + layer.w13_weight_scale, + size_k13, + layer.w13_weight_scale.shape[2], + self.group_size, + self.num_bits, + ) + replace_tensor("w13_weight_scale", marlin_w13_scales) + marlin_w2_scales = marlin_moe_permute_scales( + layer.w2_weight_scale, + layer.w2_weight_scale.shape[1] + * (self.group_size if self.group_size != -1 else self.packed_factor), + size_k2, + self.group_size, + self.num_bits, + ) + replace_tensor("w2_weight_scale", marlin_w2_scales) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + assert ( + self.moe_runner_config.activation == "silu" + ), "Only SiLU activation is supported." + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + topk_weights, topk_ids, router_logits = topk_output + + output = torch.ops.vllm.fused_marlin_moe( + x, + layer.w13_weight_packed, + layer.w2_weight_packed, + layer.w13_weight_scale, + layer.w2_weight_scale, + router_logits, + topk_weights, + topk_ids, + g_idx1=layer.w13_weight_g_idx, + g_idx2=layer.w2_weight_g_idx, + sort_indices1=layer.w13_g_idx_sort_indices, + sort_indices2=layer.w2_g_idx_sort_indices, + num_bits=self.num_bits, + is_k_full=self.is_k_full, + ) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py new file mode 100644 index 000000000..c94575316 --- /dev/null +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 + +from .compressed_tensors_scheme import CompressedTensorsScheme +from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8 +from .compressed_tensors_w8a16_fp8 import CompressedTensorsW8A16Fp8 + +__all__ = [ + "CompressedTensorsScheme", + "CompressedTensorsW8A8Fp8", + "CompressedTensorsW8A16Fp8", +] diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py new file mode 100644 index 000000000..3795d0a54 --- /dev/null +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py @@ -0,0 +1,56 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod +from typing import Optional + +import torch + +__all__ = ["CompressedTensorsScheme"] + + +class CompressedTensorsScheme(ABC): + """ + Abstract class used to describe the weight creation and forward pass + of different quantization schemes supported by CompressedTensors. + """ + + @classmethod + @abstractmethod + def get_min_capability(cls) -> int: + """ + Get minimum device capability. + """ + raise NotImplementedError + + @abstractmethod + def create_weights(self, *args, **kwargs): + """ + Weight creation for the particular scheme. Inputs to this function + + """ + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] + ): + """ + Run the forward pass for the particular scheme. This is where + scheme-specific dequant/quant steps/kernels should be applied. + + :param layer: torch.nn.Module with the registered weights and + other parameters relevant to the particular scheme. + :param x: input to the layer + :param bias: bias parameter + + """ + raise NotImplementedError + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module): + """ + Called after weight loading is complete for any cleanup that + needs to occur. + """ + raise NotImplementedError diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py new file mode 100644 index 000000000..af4f1a0e0 --- /dev/null +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py @@ -0,0 +1,153 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +from typing import Callable, List, Optional + +import torch +from compressed_tensors.quantization import QuantizationStrategy + +from sglang.srt.layers.parameter import ( + ChannelQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsScheme, +) +from sglang.srt.layers.quantization.utils import convert_to_channelwise + +try: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + prepare_fp8_layer_for_marlin, + ) + + MARLIN_FP8_AVAILABLE = True +except ImportError: + MARLIN_FP8_AVAILABLE = False + + def apply_fp8_marlin_linear(*args, **kwargs): + raise ImportError("vllm is not installed") + + def prepare_fp8_layer_for_marlin(*args, **kwargs): + raise ImportError("vllm is not installed") + + +__all__ = ["CompressedTensorsW8A16Fp8"] + +SUPPORTED_STRATEGIES = [QuantizationStrategy.CHANNEL, QuantizationStrategy.TENSOR] + + +class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): + def __init__(self, strategy: str, is_static_input_scheme: bool): + self.strategy = strategy + self.is_static_input_scheme = is_static_input_scheme + + if not MARLIN_FP8_AVAILABLE: + raise ImportError( + "vllm is not installed. To use CompressedTensorsW8A16Fp8, please install vllm" + ) + + @classmethod + def get_min_capability(cls) -> int: + # ampere and up + return 80 + + # W8A8-Fp8 kernels support only per-tensor and per-channel cases. + # So if we have a fused module (QKV, MLP) with per tensor scales, + # we expand each scale to its shard's channels. + def process_weights_after_loading(self, layer) -> None: + if self.strategy == QuantizationStrategy.TENSOR: + ws_channelwise = convert_to_channelwise( + layer.weight_scale, layer.logical_widths + ) + layer.weight_scale = torch.nn.Parameter(ws_channelwise, requires_grad=False) + else: + # required by torch.compile to be torch.nn.Parameter + layer.weight_scale = torch.nn.Parameter( + layer.weight_scale.data, requires_grad=False + ) + + # Weights must be transposed for marlin + layer.weight = torch.nn.Parameter(layer.weight.t(), requires_grad=False) + + if self.is_static_input_scheme: + # required by torch.compile to be torch.nn.Parameter + layer.input_scale = torch.nn.Parameter( + layer.input_scale.data, requires_grad=False + ) + prepare_fp8_layer_for_marlin(layer, size_k_first=True) + + def create_weights( + self, + layer: torch.nn.Module, + input_size: int, + output_partition_sizes: List[int], + input_size_per_partition: int, + params_dtype: torch.dtype, + weight_loader: Callable, + **kwargs, + ): + output_size_per_partition = sum(output_partition_sizes) + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # WEIGHT + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=torch.float8_e4m3fn, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + if self.strategy == QuantizationStrategy.CHANNEL: + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + elif self.strategy == QuantizationStrategy.TENSOR: + weight_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + else: + raise ValueError( + f"Unsupported weight strategy={self.strategy}, " + f"supported strategies are {SUPPORTED_STRATEGIES}" + ) + + weight_scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale", weight_scale) + + # INPUT SCALE (to deal with converted checkpoints) + if self.is_static_input_scheme: + input_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("input_scale", input_scale) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_fp8_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py new file mode 100644 index 000000000..a157ebc3e --- /dev/null +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py @@ -0,0 +1,172 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +from typing import Callable, List, Optional + +import torch +from compressed_tensors.quantization import QuantizationStrategy +from torch.nn import Parameter + +from sglang.srt.layers.parameter import ( + ChannelQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsScheme, +) +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.layers.quantization.utils import requantize_with_max_scale +from sglang.srt.utils import get_bool_env_var, is_hip + +__all__ = ["CompressedTensorsW8A8Fp8"] + +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +if _use_aiter: + from aiter.ops.shuffle import shuffle_weight + + +class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): + + def __init__(self, strategy: str, is_static_input_scheme: bool): + self.strategy = strategy + self.is_static_input_scheme = is_static_input_scheme + + @classmethod + def get_min_capability(cls) -> int: + # lovelace and up + return 89 + + def process_weights_after_loading(self, layer) -> None: + # If per tensor, when we have a fused module (e.g. QKV) with per + # tensor scales (thus N scales being passed to the kernel), + # requantize so we can always run per tensor + if self.strategy == QuantizationStrategy.TENSOR: + max_w_scale, weight = requantize_with_max_scale( + weight=layer.weight, + weight_scale=layer.weight_scale, + logical_widths=layer.logical_widths, + ) + + if is_fp8_fnuz(): + input_scale = getattr(layer, "input_scale", None) + + weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, weight_scale=max_w_scale, input_scale=input_scale + ) + if input_scale is not None: + layer.input_scale = Parameter(input_scale, requires_grad=False) + + layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight_scale = Parameter(max_w_scale, requires_grad=False) + + # If channelwise, scales are already lined up, so just transpose. + elif self.strategy == QuantizationStrategy.CHANNEL: + weight = layer.weight + + if is_fp8_fnuz(): + input_scale = getattr(layer, "input_scale", None) + + weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, + weight_scale=layer.weight_scale, + input_scale=input_scale, + ) + if input_scale is not None: + layer.input_scale = Parameter(input_scale, requires_grad=False) + else: + weight_scale = layer.weight_scale.data + + if _use_aiter: + layer.weight = Parameter( + shuffle_weight(weight, (16, 16)), requires_grad=False + ) + else: + layer.weight = Parameter(weight.t(), requires_grad=False) + + # required by torch.compile to be torch.nn.Parameter + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + + else: + raise ValueError(f"Unknown quantization strategy {self.strategy}") + + # INPUT SCALE + if self.is_static_input_scheme and hasattr(layer, "input_scale"): + layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) + else: + layer.input_scale = None + + def create_weights( + self, + layer: torch.nn.Module, + output_partition_sizes: List[int], + input_size_per_partition: int, + params_dtype: torch.dtype, + weight_loader: Callable, + **kwargs, + ): + output_size_per_partition = sum(output_partition_sizes) + layer.logical_widths = output_partition_sizes + + # WEIGHT + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=torch.float8_e4m3fn, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + # TODO: update create_xxx_parameter functions to return + # the newly added parameters + if self.strategy == QuantizationStrategy.CHANNEL: + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + else: + assert self.strategy == QuantizationStrategy.TENSOR + weight_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + + # min requirement for fp8 kernels + weight_scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale", weight_scale) + + # INPUT SCALE + if self.is_static_input_scheme: + input_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + input_scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("input_scale", input_scale) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_fp8_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + use_per_token_if_dynamic=True, + compressed_tensor_quant=True, + ) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/utils.py b/python/sglang/srt/layers/quantization/compressed_tensors/utils.py new file mode 100644 index 000000000..ddefa5ea3 --- /dev/null +++ b/python/sglang/srt/layers/quantization/compressed_tensors/utils.py @@ -0,0 +1,218 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +import re +from types import MappingProxyType +from typing import Iterable, List, Mapping, Optional + +from compressed_tensors import CompressionFormat +from torch.nn import Module + + +def is_activation_quantization_format(format: str) -> bool: + _ACTIVATION_QUANTIZATION_FORMATS = [ + CompressionFormat.naive_quantized.value, + CompressionFormat.int_quantized.value, + CompressionFormat.float_quantized.value, + ] + return format in _ACTIVATION_QUANTIZATION_FORMATS + + +def should_ignore_layer( + layer_name: Optional[str], + ignore: Iterable[str] = tuple(), + fused_mapping: Mapping[str, List[str]] = MappingProxyType({}), +) -> bool: + if layer_name is None: + return False + + # layer_name = model.layers.0.self_attn.qkv_proj + # proj_name = qkv_proj + proj_name = layer_name.split(".")[-1] + + # Fused layers like gate_up_proj or qkv_proj will not be fused + # in the safetensors checkpoint. So, we convert the name + # from the fused version to unfused + check to make sure that + # each shard of the fused layer has the same scheme. + if proj_name in fused_mapping and layer_name not in ignore: + shard_proj_names = fused_mapping[proj_name] + + # Convert fused_name --> [shard_names] + shard_names = [ + layer_name.replace(proj_name, shard_proj_name) + for shard_proj_name in shard_proj_names + ] + + # Layer should be ignored if shards are ignored. + should_ignore_layer = None + for shard_name in shard_names: + should_ignore_shard = check_equal_or_regex_match( + layer_name=shard_name, targets=ignore + ) + + # If shard_idx=0, set layer ignore to match shard. + if should_ignore_layer is None: + should_ignore_layer = should_ignore_shard + + # If shard_idx=1+ confirm scheme matches prior shards. + elif should_ignore_shard != should_ignore_layer: + raise ValueError( + f"Found a different quantization schemes for " + f"{shard_proj_names} in {layer_name}. vLLM " + "requires all to use the same scheme." + ) + + # Unfused layers like down_proj and o_proj will match + # the safetensors checkpoint already. + else: + should_ignore_layer = check_equal_or_regex_match( + layer_name=layer_name, targets=ignore + ) + + assert should_ignore_layer is not None + return should_ignore_layer + + +def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool: + """ + Checks whether a layer_name is exactly equal or a regex match for + if target starts with 're:' to any target in list. + """ + for target in targets: + if _is_equal_or_regex_match(layer_name, target): + return True + return False + + +def find_matched_target( + layer_name: Optional[str], + module: Module, + targets: Iterable[str], + fused_mapping: Mapping[str, List[str]] = MappingProxyType({}), +) -> str: + """ + Helper function to look up which "target" in the compressed-tensors + config that a layer corresponds to. + + Recall that a compressed-tensors configs has a concept of + config_groups, where each layer can be quantized with with a different + scheme. + + targets in each config_group will be a list of either layer names + (or regexes corresponding to layer names) or names of torch Modules. + + First, we try to match the layer_name with a target + Second, we try to match the module's name with a target + Third, we try to map the layer_name to a list of fused module names. + *All* component module names must match in order for a match to be + successful. A successful match returns the first component target + + :param layer_name: layer name + :param module: torch.nn.Module + :param targets: list of targets to match the layer against + :param fused_mapping: map from fused layer names to its components + :param fused_strategy: either "all" or "any". If using "all", fused + layers match if "all" of its components match + """ + + if layer_name is None: + layer_name = "" + + matched_target = ( + _find_first_match(layer_name, targets) + or _find_first_match(module.__class__.__name__, targets, True) + or _match_fused_layer(layer_name, targets, fused_mapping) + ) + + if matched_target is None: + raise ValueError( + f"Unable to find matching target for {layer_name} in the " + "compressed-tensors config." + ) + + return matched_target + + +def _find_first_match( + value: str, targets: Iterable[str], check_contains: bool = False +) -> Optional[str]: + """ + Returns first element of target that matches value either + exactly or as a regex after 're:'. If check_contains is set to True, + additionally checks if the target string is contained within the value. + + :param value: string to compare the list of targets against + :param targets: list of targets to match the layer against + :param check_contains: whether or not to do a substring match + """ + + for target in targets: + if _is_equal_or_regex_match(value, target, check_contains=check_contains): + return target + return None + + +def _is_equal_or_regex_match( + value: str, target: str, check_contains: bool = False +) -> bool: + """ + Checks whether a value is exactly equal or a regex match for target + if target starts with 're:'. If check_contains is set to True, + additionally checks if the target string is contained within the value. + """ + + if target.startswith("re:"): + pattern = target[3:] + if re.match(pattern, value): + return True + elif check_contains: + if target.lower() in value.lower(): + return True + elif target == value: + return True + return False + + +def _match_fused_layer( + layer_name: str, + target_layers: Iterable[str], + fused_mapping: Mapping[str, List[str]], +) -> Optional[str]: + """ + Match a fused layer name to its corresponding individual layer in + target_layers. Returns first value in fused_mapping which matches targets + + Implements an "all" matching strategy where a fused layer matches iff + "all" of its components match + + :param layer_name: layer name + :param target_layers: list of targets to match the layer against + :param fused_mapping: map from fused layer names to its components + + Examples: + layer_name = "model.layers.0.self_attn.qkv_proj" + target_layers = ["model.layers.0.self_attn.q_proj", + "model.layers.0.self_attn.k_proj", + "model.layers.0.self_attn.v_proj"] + """ + # find layer_name in mapping + fused = next((key for key in fused_mapping if layer_name.endswith(key)), None) + if fused is None: + return None + + # expand path of unfused components + unfused_paths = [ + layer_name.replace(fused, unfused) for unfused in fused_mapping[fused] + ] + + # for each unfused component, find a match in targets + unfused_matches: List[Optional[str]] = [] + for unfused in unfused_paths: + for target in target_layers: + if _is_equal_or_regex_match(unfused, target): + unfused_matches.append(target) + break + else: + unfused_matches.append(None) + + return unfused_matches[0] if all(unfused_matches) else None diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6ae89c757 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..ae4c397ef --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6496a38fb --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..effb7bda7 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..eb1848894 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c098ef2db --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c098ef2db --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..3633e363f --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..653bb997e --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..31d29e10c --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..3618053b6 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..cb573ccdf --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..46a982f5e --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a5679be13 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..7d8044772 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..eb11c1696 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..035ec027f --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..62ccf8e6f --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..8b49f2781 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..dd66019cc --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5011fd572 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..1957a7f32 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..851bc9f9f --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6c2eebe0f --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..d1227c215 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..448daf360 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6f5adbb93 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6f5adbb93 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6f5adbb93 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..7b10d6d78 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..e76d41dfd --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..77ba0d747 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..1c61451fb --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..66ca4ee16 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..48b64a87d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..63e661c80 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..d521e2bd0 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..d72288781 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c65a16f3d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..cf3540379 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a8ab31bfe --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..1477e063d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6f5adbb93 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6f5adbb93 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..80d23e969 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..0a5d7bfdb --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..970345270 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..eccb86a76 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..cb91a279d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..88af48431 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a794fc821 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..dd069726d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..4225c78eb --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..4225c78eb --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..4225c78eb --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..933170b7b --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..797ee3d45 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..7febe3d27 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..56b939e52 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..1c36fe951 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..30932a1cc --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..63d9a0bf5 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..65989f012 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=36864,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=36864,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..7fa398c15 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=36864,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=36864,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=36864,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..f15d8f64c --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=36864,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..4225c78eb --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..4225c78eb --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a5518d979 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..9d7658bfc --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..cd3e07804 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..932321cca --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..9d5a329d7 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5e6789d00 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5e6789d00 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..44c67c008 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..03dba5ad1 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..96e1594a3 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..017dca003 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5ffd367df --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..49ac14d2a --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..49ac14d2a --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..09502f05c --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..9a5ff48b8 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..b57588f9b --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..eabc42394 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..dcbb0efc5 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..dcbb0efc5 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..7fd942e3e --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..13e3c20b7 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..ecbca9d8d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..386928de1 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..51e237b91 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a8ab31bfe --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..22ca71e2c --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6280219c9 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..78db73078 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5e40bbe64 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..044868d2d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..40c01c0b9 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..b96981dec --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c6fd36597 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..f4e026497 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..246062a22 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..0b24efa37 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..160f12ed3 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a3a434e12 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..e5c4a1d2c --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..2f3962f68 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..bfdff8adc --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..d694716c3 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..2bf5eb27e --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..871a0cf2e --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..f211a5d38 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..dfe5c1e43 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..dfe5c1e43 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..dfe5c1e43 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..b2cab0831 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..9de2ba9c8 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..9c908e804 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..0a1e14cff --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c4b1bd435 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..c87cb4f95 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..15b1c93f6 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5bf208ce4 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..6ae00329e --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..59b79432f --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..f78e7060e --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..8ff12e64c --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..e997126f9 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..b474e6265 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..4532f9368 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a3f5371f1 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..dfe5c1e43 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..dfe5c1e43 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..72ab3dff2 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..1d3ce5c94 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..ca7f32b95 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..12e305853 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5acea242c --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a87f5de1b --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a87f5de1b --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..37208577d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..3ab5796ee --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..58cdd93e9 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..a5ce5e49b --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..b72e0371d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..468f9e78d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..468f9e78d --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..d8cc0f896 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..3cb7eaa07 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..129f59e0c --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..293adce38 --- /dev/null +++ b/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/python/sglang/srt/layers/quantization/deep_gemm_wrapper/__init__.py b/python/sglang/srt/layers/quantization/deep_gemm_wrapper/__init__.py new file mode 100644 index 000000000..32f3ebe04 --- /dev/null +++ b/python/sglang/srt/layers/quantization/deep_gemm_wrapper/__init__.py @@ -0,0 +1 @@ +from .entrypoint import * diff --git a/python/sglang/srt/layers/quantization/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/quantization/deep_gemm_wrapper/compile_utils.py new file mode 100644 index 000000000..e374759c4 --- /dev/null +++ b/python/sglang/srt/layers/quantization/deep_gemm_wrapper/compile_utils.py @@ -0,0 +1,241 @@ +import logging +import os +from contextlib import contextmanager +from enum import IntEnum, auto +from typing import Dict, List, Tuple + +import torch +from tqdm import tqdm + +from sglang.srt.layers.quantization.deep_gemm_wrapper.configurer import ( + ENABLE_JIT_DEEPGEMM, +) +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import ceil_div, get_bool_env_var, get_int_env_var + +logger = logging.getLogger(__name__) + +if ENABLE_JIT_DEEPGEMM: + import deep_gemm + + +_BUILTIN_M_LIST = list(range(1, 1024 * 16 + 1)) +_ENABLE_JIT_DEEPGEMM_PRECOMPILE = get_bool_env_var( + "SGL_JIT_DEEPGEMM_PRECOMPILE", "true" +) +_DO_COMPILE_ALL = True +_IS_FIRST_RANK_ON_NODE = get_bool_env_var("SGL_IS_FIRST_RANK_ON_NODE", "true") +_COMPILE_WORKERS = get_int_env_var("SGL_JIT_DEEPGEMM_COMPILE_WORKERS", 4) +_IN_PRECOMPILE_STAGE = get_bool_env_var("SGL_IN_DEEPGEMM_PRECOMPILE_STAGE", "false") + +# Force redirect deep_gemm cache_dir +os.environ["DG_JIT_CACHE_DIR"] = os.getenv( + "SGL_DG_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm") +) + +# Refer to https://github.com/deepseek-ai/DeepGEMM/commit/d75b218b7b8f4a5dd5406ac87905039ead3ae42f +# NVRTC may have performance loss with some cases. +# And NVCC JIT speed is also 9x faster in the ref commit +os.environ["DG_JIT_USE_NVRTC"] = os.getenv("SGL_DG_USE_NVRTC", "0") + + +def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): + global _BUILTIN_M_LIST + global _DO_COMPILE_ALL + global _IS_FIRST_RANK_ON_NODE + + # Generate m_max + m_max = 1024 * 16 + if server_args.chunked_prefill_size < 1: + m_max = 1024 * 64 + elif server_args.chunked_prefill_size > 8192: + m_max = server_args.chunked_prefill_size * 2 + m_max = min(1024 * 128, m_max) + _BUILTIN_M_LIST = list(range(1, m_max + 1)) + + _IS_FIRST_RANK_ON_NODE = ServerArgs.base_gpu_id == gpu_id + + # Check if is the first rank on node. + # Default each rank will try compile all Ms to + # load all symbols at the launch stages. + # Avoid loading symbols at the serving stages. + _DO_COMPILE_ALL = _IS_FIRST_RANK_ON_NODE + + +class DeepGemmKernelType(IntEnum): + GROUPED_GEMM_NT_F8F8BF16_MASKED = auto() + GROUPED_GEMM_NT_F8F8BF16_CONTIG = auto() + GEMM_NT_F8F8BF16 = auto() + + +_INITIALIZATION_DICT: Dict[Tuple[DeepGemmKernelType, int, int, int], bool] = dict() + + +# TODO improve code +def _maybe_compile_deep_gemm_one_type_all( + kernel_type: DeepGemmKernelType, + n: int, + k: int, + num_groups: int, +) -> None: + global _INITIALIZATION_DICT + global _BUILTIN_M_LIST + + query_key = (kernel_type, n, k, num_groups) + if ( + _ENABLE_JIT_DEEPGEMM_PRECOMPILE + and _DO_COMPILE_ALL + and _INITIALIZATION_DICT.get(query_key) is None + ): + _INITIALIZATION_DICT[query_key] = True + + # TODO maybe improve logs + if not _IN_PRECOMPILE_STAGE and _IS_FIRST_RANK_ON_NODE: + logger.warning( + "Entering DeepGEMM JIT Pre-Compile session. " + "It may take a long time (typically 10-20 mins) " + "if you have not run `sglang.compile_deep_gemm`. " + "It is recommended to run `sglang.compile_deep_gemm` with same args as `sglang.launch_server`" + " for pre-compilation to reduce the overhead if you have not run it before. " + "For example: " + "`python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code`" + ) + + logger.info( + f"Try DeepGEMM JIT Compiling for " + f"<{kernel_type.name}> N={n}, K={k}, num_groups={num_groups} with all Ms." + f"{' It only takes a little time (typically 1 sec) if you have run `python3 -m sglang.compile_deep_gemm`. ' if not _IN_PRECOMPILE_STAGE else ''}" + ) + + _compile_deep_gemm_one_type_all( + kernel_type=kernel_type, + n=n, + k=k, + num_groups=num_groups, + m_list=_BUILTIN_M_LIST, + ) + + +# NOTE(alcanderian): get_num_sms should be change when 2-batch-overlap is introduced +def _compile_deep_gemm_one_type_all( + kernel_type: DeepGemmKernelType, + n: int, + k: int, + num_groups: int, + m_list: List[int], +) -> None: + if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: + m_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout() + m_list = sorted(list(set(m for m in m_list if m % m_alignment == 0))) + + executor = _BaseWarmupExecutor.create( + kernel_type, max_m=max(m_list), n=n, k=k, num_groups=num_groups + ) + + old_compile_mode = deep_gemm.get_compile_mode() + deep_gemm.set_compile_mode(1) + # TODO can use multi thread + for m in tqdm(m_list, desc=f"DeepGEMM warmup"): + executor.execute(m=m) + deep_gemm.set_compile_mode(old_compile_mode) + + # clean up input buffers + torch.cuda.current_stream().synchronize() + del executor + torch.cuda.empty_cache() + + +class _BaseWarmupExecutor: + @staticmethod + def create(kernel_type: DeepGemmKernelType, **kwargs): + return { + DeepGemmKernelType.GEMM_NT_F8F8BF16: _NormalWarmupExecutor, + DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: _GroupedContWarmupExecutor, + DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED: _GroupedMaskedWarmupExecutor, + }[kernel_type](**kwargs) + + def execute(self, m): + raise NotImplementedError + + +def _empty_token_fp8(size): + *dims, k = size + return ( + torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn), + torch.empty( + (*dims, ceil_div(k, _BLOCK_SIZE)), device="cuda", dtype=torch.float32 + ), + ) + + +def _empty_block_fp8(size): + *dims, n, k = size + return ( + torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn), + torch.empty( + (*dims, ceil_div(n, _BLOCK_SIZE), ceil_div(k, _BLOCK_SIZE)), + device="cuda", + dtype=torch.float32, + ), + ) + + +_BLOCK_SIZE = 128 + + +class _NormalWarmupExecutor(_BaseWarmupExecutor): + def __init__(self, max_m: int, n: int, k: int, num_groups: int): + self.lhs_q, self.lhs_s = _empty_token_fp8((max_m, k)) + self.rhs_q, self.rhs_s = _empty_block_fp8((n, k)) + self.out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16) + + def execute(self, m): + deep_gemm.fp8_gemm_nt( + (self.lhs_q[:m], self.lhs_s[:m]), + (self.rhs_q, self.rhs_s), + self.out[:m], + ) + + +class _GroupedContWarmupExecutor(_BaseWarmupExecutor): + def __init__(self, max_m: int, n: int, k: int, num_groups: int): + self.lhs_q, self.lhs_s = _empty_token_fp8((max_m, k)) + self.rhs_q, self.rhs_s = _empty_block_fp8((num_groups, n, k)) + self.m_indices = torch.zeros((max_m,), device="cuda", dtype=torch.int32) + self.out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16) + + def execute(self, m): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + (self.lhs_q[:m], self.lhs_s[:m]), + (self.rhs_q, self.rhs_s), + self.out[:m], + m_indices=self.m_indices[:m], + ) + + +class _GroupedMaskedWarmupExecutor(_BaseWarmupExecutor): + def __init__(self, max_m: int, n: int, k: int, num_groups: int): + self.lhs_q, self.lhs_s = _empty_token_fp8((num_groups, max_m, k)) + self.rhs_q, self.rhs_s = _empty_block_fp8((num_groups, n, k)) + self.masked_m = torch.zeros((num_groups,), device="cuda", dtype=torch.int32) + self.out = torch.empty( + (num_groups, max_m, n), device="cuda", dtype=torch.bfloat16 + ) + + def execute(self, m): + deep_gemm.fp8_m_grouped_gemm_nt_masked( + (self.lhs_q, self.lhs_s), + (self.rhs_q, self.rhs_s), + self.out, + masked_m=self.masked_m, + # DeepGEMM uses `expect_m` instead of input shape for `get_best_config` + expected_m=m, + ) + + +@contextmanager +def deep_gemm_execution_hook( + m: int, n: int, k: int, num_groups: int, kernel_type: DeepGemmKernelType +): + _maybe_compile_deep_gemm_one_type_all(kernel_type, n, k, num_groups) + yield diff --git a/python/sglang/srt/layers/quantization/deep_gemm_wrapper/configurer.py b/python/sglang/srt/layers/quantization/deep_gemm_wrapper/configurer.py new file mode 100644 index 000000000..ecf7d1647 --- /dev/null +++ b/python/sglang/srt/layers/quantization/deep_gemm_wrapper/configurer.py @@ -0,0 +1,32 @@ +import logging + +import torch + +from sglang.srt.utils import get_bool_env_var, get_device_sm + +logger = logging.getLogger(__name__) + + +def _compute_enable_deep_gemm(): + sm_version = get_device_sm() + if sm_version < 90: + return False + + try: + import deep_gemm + except ImportError: + logger.warning("Failed to import deep_gemm, disable ENABLE_JIT_DEEPGEMM.") + return False + + return get_bool_env_var("SGL_ENABLE_JIT_DEEPGEMM", default="true") + + +def _is_blackwell_arch() -> bool: + major, minor = torch.cuda.get_device_capability(torch.cuda.current_device()) + return major == 10 + + +ENABLE_JIT_DEEPGEMM = _compute_enable_deep_gemm() + +DEEPGEMM_BLACKWELL = ENABLE_JIT_DEEPGEMM and _is_blackwell_arch() +DEEPGEMM_SCALE_UE8M0 = DEEPGEMM_BLACKWELL diff --git a/python/sglang/srt/layers/quantization/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/quantization/deep_gemm_wrapper/entrypoint.py new file mode 100644 index 000000000..02945f449 --- /dev/null +++ b/python/sglang/srt/layers/quantization/deep_gemm_wrapper/entrypoint.py @@ -0,0 +1,119 @@ +import logging +from contextlib import contextmanager +from typing import Tuple + +import torch + +from sglang.srt.layers.quantization.deep_gemm_wrapper import compile_utils +from sglang.srt.layers.quantization.deep_gemm_wrapper.configurer import ( + DEEPGEMM_BLACKWELL, + DEEPGEMM_SCALE_UE8M0, + ENABLE_JIT_DEEPGEMM, +) +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import get_bool_env_var + +logger = logging.getLogger(__name__) + +if ENABLE_JIT_DEEPGEMM: + import deep_gemm + from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor + +_SANITY_CHECK = get_bool_env_var("SGLANG_DEEPGEMM_SANITY_CHECK") + + +# TODO maybe rename these functions +def grouped_gemm_nt_f8f8bf16_masked( + lhs: Tuple[torch.Tensor, torch.Tensor], + rhs: Tuple[torch.Tensor, torch.Tensor], + out: torch.Tensor, + masked_m: torch.Tensor, + expected_m: int, +): + num_groups, _, k = lhs[0].shape + _, n, _ = rhs[0].shape + kernel_type = compile_utils.DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED + + _sanity_check_input(lhs) + _sanity_check_input(rhs) + + with compile_utils.deep_gemm_execution_hook( + expected_m, n, k, num_groups, kernel_type + ): + deep_gemm.fp8_m_grouped_gemm_nt_masked( + lhs, + rhs, + out, + masked_m, + expected_m, + ) + + +def grouped_gemm_nt_f8f8bf16_contig( + lhs: Tuple[torch.Tensor, torch.Tensor], + rhs: Tuple[torch.Tensor, torch.Tensor], + out: torch.Tensor, + m_indices: torch.Tensor, +): + m, k = lhs[0].shape + num_groups, n, _ = rhs[0].shape + kernel_type = compile_utils.DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG + + _sanity_check_input(lhs) + _sanity_check_input(rhs) + + with compile_utils.deep_gemm_execution_hook(m, n, k, num_groups, kernel_type): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous(lhs, rhs, out, m_indices) + + +def gemm_nt_f8f8bf16( + lhs: Tuple[torch.Tensor, torch.Tensor], + rhs: Tuple[torch.Tensor, torch.Tensor], + out: torch.Tensor, +): + m, k = lhs[0].shape + n, _ = rhs[0].shape + num_groups = 1 + kernel_type = compile_utils.DeepGemmKernelType.GEMM_NT_F8F8BF16 + + _sanity_check_input(lhs) + _sanity_check_input(rhs) + + with compile_utils.deep_gemm_execution_hook(m, n, k, num_groups, kernel_type): + deep_gemm.fp8_gemm_nt( + lhs, + rhs, + out, + ) + + +def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): + compile_utils.update_deep_gemm_config(gpu_id, server_args) + + +@contextmanager +def configure_deep_gemm_num_sms(num_sms): + if num_sms is None: + yield + else: + original_num_sms = deep_gemm.get_num_sms() + deep_gemm.set_num_sms(num_sms) + try: + yield + finally: + deep_gemm.set_num_sms(original_num_sms) + + +def _sanity_check_input(x_fp8: Tuple[torch.Tensor, torch.Tensor]): + if not _SANITY_CHECK: + return + + x, x_scale = x_fp8 + + if x_scale.dtype == torch.int: + return + + from sglang.srt.layers.quantization.fp8_utils import ceil_to_ue8m0 + + x_scale_ceil = ceil_to_ue8m0(x_scale) + assert torch.all(x_scale == x_scale_ceil), f"{x_scale=} {x_scale_ceil=}" diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py new file mode 100644 index 000000000..b020e4188 --- /dev/null +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -0,0 +1,1240 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/layers/quantization/fp8.py + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import torch +import torch.nn.functional as F +from torch.nn import Module +from torch.nn.parameter import Parameter + +try: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + prepare_fp8_layer_for_marlin, + ) + + MARLIN_FP8_AVAILABLE = True +except ImportError: + MARLIN_FP8_AVAILABLE = False + + def dummy_func(*args, **kwargs): + raise ImportError( + "marlin FP8 requires some operators from vllm. Please install vllm." + ) + + apply_fp8_marlin_linear = prepare_fp8_layer_for_marlin = dummy_func + + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker +from sglang.srt.layers.parameter import ( + BlockQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8_kernel import ( + fp8_dtype, + is_fp8_fnuz, + per_token_group_quant_fp8, + scaled_fp8_quant, +) +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + can_auto_enable_marlin_fp8, + cutlass_fp8_supported, + dispatch_w8a8_block_fp8_linear, + input_to_float8, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.layers.quantization.utils import ( + all_close_1d, + convert_to_channelwise, + is_layer_skipped, + per_tensor_dequantize, + requantize_with_max_scale, +) +from sglang.srt.utils import ( + cpu_has_amx_support, + get_bool_env_var, + is_cpu, + is_cuda, + is_hip, + is_npu, + is_sm90_supported, + is_sm100_supported, + log_info_on_rank0, + next_power_of_2, + print_warning_once, + set_weight_attrs, + use_intel_amx_backend, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + DispatchOutput, + StandardDispatchOutput, + ) + from sglang.srt.layers.moe.topk import TopKOutput + from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config + +_is_hip = is_hip() +_is_cuda = is_cuda() +_is_npu = is_npu() +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() + +_is_fp8_fnuz = is_fp8_fnuz() + +_use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _is_hip and (_use_aiter or _use_hip_int4): + from aiter import ActivationType, QuantType + from aiter.fused_moe import fused_moe + from aiter.ops.shuffle import shuffle_weight + + +ACTIVATION_SCHEMES = ["static", "dynamic"] + +logger = logging.getLogger(__name__) + + +class Fp8Config(QuantizationConfig): + """Config class for FP8.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = False, + activation_scheme: str = "dynamic", + ignored_layers: Optional[List[str]] = None, + weight_block_size: List[int] = None, + ) -> None: + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + if is_checkpoint_fp8_serialized: + log_info_on_rank0(logger, "Detected fp8 checkpoint.") + if activation_scheme not in ACTIVATION_SCHEMES: + raise ValueError(f"Unsupported activation scheme {activation_scheme}") + self.activation_scheme = activation_scheme + self.ignored_layers = ignored_layers or [] + if weight_block_size is not None: + if not is_checkpoint_fp8_serialized: + raise ValueError( + f"The block-wise quantization only supports fp8-serialized checkpoint for now." + ) + if len(weight_block_size) != 2: + raise ValueError( + f"The quantization block size of weight must have 2 dimensions, but got {len(weight_block_size)} dimensions." + ) + if activation_scheme != "dynamic": + raise ValueError( + f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme." + ) + self.weight_block_size = weight_block_size + + @classmethod + def get_name(cls) -> str: + return "fp8" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> Fp8Config: + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_fp8_serialized = "fp8" in quant_method + activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) + return cls( + is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + weight_block_size=weight_block_size, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + if is_layer_skipped(prefix, self.ignored_layers): + return UnquantizedLinearMethod() + return Fp8LinearMethod(self) + elif isinstance(layer, FusedMoE): + return Fp8MoEMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class Fp8LinearMethod(LinearMethodBase): + """Linear method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + + Limitations: + 1. Only support per-tensor quantization due to torch._scaled_mm support. + 2. Only support float8_e4m3fn data type due to the limitation of + torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856) + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Union[Fp8Config, W4AFp8Config]): + self.quant_config = quant_config + self.cutlass_fp8_supported = cutlass_fp8_supported() + + # For GPUs that lack FP8 hardware support, we can leverage the Marlin + # kernel for fast weight-only FP8 quantization + self.use_marlin = False + if _is_cuda and MARLIN_FP8_AVAILABLE: + force_marlin = get_bool_env_var("SGLANG_FORCE_FP8_MARLIN") + auto_enable = can_auto_enable_marlin_fp8() + self.use_marlin = force_marlin or auto_enable + + self.block_quant = self.quant_config.weight_block_size is not None + + self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + tp_size = get_tensor_model_parallel_world_size() + if self.block_quant: + block_n, block_k = ( + self.quant_config.weight_block_size[0], + self.quant_config.weight_block_size[1], + ) + # Required by row parallel + if tp_size > 1 and input_size // input_size_per_partition == tp_size: + if input_size_per_partition % block_k != 0: + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + # Required by column parallel or enabling merged weights + if ( + tp_size > 1 and output_size // output_size_per_partition == tp_size + ) or len(output_partition_sizes) > 1: + for output_partition_size in output_partition_sizes: + if output_partition_size % block_n != 0: + raise ValueError( + f"Weight output_partition_size = " + f"{output_partition_size} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # WEIGHT + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_fp8_serialized + else params_dtype + ) + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, input_size_per_partition, dtype=weight_dtype + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # If checkpoint is serialized fp8, load them. + # Otherwise, wait until process_weights_after_loading. + if self.quant_config.is_checkpoint_fp8_serialized: + # WEIGHT SCALE + if self.block_quant: + if hasattr(self.quant_config, "activation_scheme"): + assert self.quant_config.activation_scheme == "dynamic" + elif hasattr(self.quant_config, "linear_activation_scheme"): + assert self.quant_config.linear_activation_scheme == "dynamic" + scale = BlockQuantScaleParameter( + data=torch.empty( + (output_size_per_partition + block_n - 1) // block_n, + (input_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale_inv", scale) + else: + scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale", scale) + + # INPUT ACTIVATION SCALE + if ( + hasattr(self.quant_config, "activation_scheme") + and self.quant_config.activation_scheme == "static" + ) or ( + hasattr(self.quant_config, "linear_activation_scheme") + and self.quant_config.linear_activation_scheme == "static" + ): + scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("input_scale", scale) + else: + layer.register_parameter("input_scale", None) + + def process_weights_after_loading(self, layer: Module) -> None: + if self.block_quant: + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_fp8_fnuz: + # activation_scheme: dynamic + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.weight, + weight_scale=layer.weight_scale_inv, + input_scale=None, + ) + layer.input_scale = None + elif _is_cpu: + assert ( + _is_cpu_amx_available + ), "Fp8LinearMethod on CPU requires that CPU has AMX support" + _amx_process_weight_after_loading(layer, ["weight"]) + layer.weight_scale_inv = torch.nn.Parameter( + layer.weight_scale_inv.data, requires_grad=False + ) + return + else: + weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data + layer.weight = Parameter(weight, requires_grad=False) + layer.weight_scale_inv = Parameter(weight_scale, requires_grad=False) + else: + layer.weight = Parameter(layer.weight.data, requires_grad=False) + + # If checkpoint not serialized fp8, quantize the weights. + if not self.quant_config.is_checkpoint_fp8_serialized: + if self.cutlass_fp8_supported or self.use_marlin: + # apply per-channel quantization default as + # cutlass sgl-kernel and marlin only support per-channel scale + qweight, weight_scale = per_token_group_quant_fp8( + layer.weight, layer.weight.shape[-1] + ) + weight_scale = weight_scale.t().contiguous() + else: + # per-tensor quantization + qweight, weight_scale = input_to_float8(layer.weight) + + # Update the layer with the new values. + layer.weight = Parameter(qweight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + layer.input_scale = None + + # If checkpoint is fp8, handle that there are N scales for N + # shards in a fused module + else: + layer.weight_scale = Parameter( + layer.weight_scale.data, requires_grad=False + ) + if ( + hasattr(self.quant_config, "activation_scheme") + and self.quant_config.activation_scheme == "static" + ) or ( + hasattr(self.quant_config, "linear_activation_scheme") + and self.quant_config.linear_activation_scheme == "static" + ): + layer.input_scale = Parameter( + layer.input_scale.data, requires_grad=False + ) + + # cutlass sgl-kernel and marlin only support per-channel scale + if self.cutlass_fp8_supported or self.use_marlin: + weight = layer.weight + weight_scale = convert_to_channelwise( + layer.weight_scale, layer.logical_widths + ) + else: + # Dequant -> Quant with max scale so we can run per tensor. + weight = layer.weight + weight_scale = layer.weight_scale + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_fp8_fnuz: + weight, weight_scale, input_scale = ( + normalize_e4m3fn_to_e4m3fnuz( + weight=weight, + weight_scale=weight_scale, + input_scale=layer.input_scale, + ) + ) + if input_scale is not None: + layer.input_scale = Parameter( + input_scale, requires_grad=False + ) + + weight_scale, weight = requantize_with_max_scale( + weight=weight, + weight_scale=weight_scale, + logical_widths=layer.logical_widths, + ) + + # Update layer with new values. + layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + if ( + hasattr(self.quant_config, "activation_scheme") + and self.quant_config.activation_scheme == "static" + ) or ( + hasattr(self.quant_config, "linear_activation_scheme") + and self.quant_config.linear_activation_scheme == "static" + ): + layer.input_scale = Parameter( + layer.input_scale.max(), requires_grad=False + ) + + if self.use_marlin: + if self.block_quant: + layer.weight_block_size = self.quant_config.weight_block_size + prepare_fp8_layer_for_marlin(layer, not self.block_quant) + # Activations not quantized for marlin. + del layer.input_scale + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if self.use_marlin: + return apply_fp8_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) + + if self.block_quant: + if use_intel_amx_backend(layer): + return torch.ops.sgl_kernel.fp8_scaled_mm_cpu( + x, + layer.weight, + layer.weight_scale_inv, + self.quant_config.weight_block_size, + bias, + x.dtype, + True, # is_vnni + ) + + return self.w8a8_block_fp8_linear( + input=x, + weight=layer.weight, + block_size=self.quant_config.weight_block_size, + weight_scale=layer.weight_scale_inv, + input_scale=None, + bias=bias, + ) + + return apply_fp8_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + use_per_token_if_dynamic=False, + ) + + +def get_tile_tokens_dim(num_tokens, top_k, num_experts): + # Guess tokens per expert assuming perfect expert distribution first. + num_tokens_per_expert = (num_tokens * top_k) // num_experts + # And pad the number to the next power of 2. + tile_tokens_dim = next_power_of_2(num_tokens_per_expert) + # Cap to 8-64 tokens per CTA tile as it's the range supported by the kernel. + tile_tokens_dim = min(max(tile_tokens_dim, 8), 64) + return tile_tokens_dim + + +class Fp8MoEMethod(FusedMoEMethodBase): + """MoE method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Fp8Config): + self.quant_config = quant_config + self.block_quant = self.quant_config.weight_block_size is not None + self.cutlass_fp8_supported = cutlass_fp8_supported() + self.use_cutlass_fused_experts_fp8 = ( + get_bool_env_var("SGLANG_CUTLASS_MOE") + and self.cutlass_fp8_supported + and self.block_quant + and (is_sm100_supported() or is_sm90_supported()) + ) + + def create_weights( + self, + layer: Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + if self.quant_config.is_checkpoint_fp8_serialized: + params_dtype = torch.uint32 if _use_hip_int4 else torch.float8_e4m3fn + tp_size = get_tensor_model_parallel_world_size() + if self.block_quant: + block_n, block_k = ( + self.quant_config.weight_block_size[0], + self.quant_config.weight_block_size[1], + ) + # NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n. + # Required by column parallel or enabling merged weights + if intermediate_size_per_partition % block_n != 0: + raise ValueError( + f"The output_size of gate's and up's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + if tp_size > 1: + # Required by row parallel + if intermediate_size_per_partition % block_k != 0: + raise ValueError( + f"The input_size of down's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + + # WEIGHTS + if _is_hip and _use_hip_int4: + # INT4 MoE weight - INT32 packed + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // 8, + dtype=params_dtype, + ), + requires_grad=False, + ) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // 8, + dtype=params_dtype, + ), + requires_grad=False, + ) + else: + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + if self.block_quant: + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * ((intermediate_size_per_partition + block_n - 1) // block_n), + (hidden_size + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + (hidden_size + block_n - 1) // block_n, + (intermediate_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale_inv", w13_weight_scale) + layer.register_parameter("w2_weight_scale_inv", w2_weight_scale) + assert self.quant_config.activation_scheme == "dynamic" + if self.use_cutlass_fused_experts_fp8: + self.ab_strides1 = torch.full( + (num_experts,), + hidden_size, + device=w13_weight.device, + dtype=torch.int64, + ) + self.c_strides1 = torch.full( + (num_experts,), + 2 * intermediate_size_per_partition, + device=w13_weight.device, + dtype=torch.int64, + ) + self.ab_strides2 = torch.full( + (num_experts,), + intermediate_size_per_partition, + device=w2_weight.device, + dtype=torch.int64, + ) + self.c_strides2 = torch.full( + (num_experts,), + hidden_size, + device=w2_weight.device, + dtype=torch.int64, + ) + self.workspace = torch.empty( + 90000, device=w13_weight.device, dtype=torch.uint8 + ) + self.a_ptr = torch.empty( + num_experts, device=w13_weight.device, dtype=torch.int64 + ) + self.b_ptr = torch.empty( + num_experts, device=w13_weight.device, dtype=torch.int64 + ) + self.out_ptr = torch.empty( + num_experts, device=w13_weight.device, dtype=torch.int64 + ) + self.a_scales_ptr = torch.empty( + num_experts, device=w13_weight.device, dtype=torch.int64 + ) + self.b_scales_ptr = torch.empty( + num_experts, device=w13_weight.device, dtype=torch.int64 + ) + self.expert_offsets = torch.empty( + num_experts + 1, device=w13_weight.device, dtype=torch.int32 + ) + self.problem_sizes1 = torch.empty( + num_experts, 3, device=w13_weight.device, dtype=torch.int32 + ) + self.problem_sizes2 = torch.empty( + num_experts, 3, device=w13_weight.device, dtype=torch.int32 + ) + + else: + # Allocate 2 scales for w1 and w3 respectively. + # They will be combined to a single scale after weight loading. + w13_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + if _is_hip: # _use_aiter: TODO: add check back after triton kernel + # ROCm - using column scaling, duplicate scaling numbers in case per tensor scaling + w13_weight_scale1 = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * intermediate_size_per_partition, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale1 = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale1", w13_weight_scale1) + layer.register_parameter("w2_weight_scale1", w2_weight_scale1) + + # Add the quantization method used (per tensor/grouped/channel) + # to ensure the weight scales are loaded in properly + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + if self.block_quant + else {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + # If loading fp8 checkpoint, pass the weight loaders. + # If loading an fp16 checkpoint, do not (we will quantize in + # process_weights_after_loading() + if self.quant_config.is_checkpoint_fp8_serialized: + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + if _is_hip and _use_hip_int4: + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + set_weight_attrs(w13_weight_scale1, extra_weight_attrs) + set_weight_attrs(w2_weight_scale1, extra_weight_attrs) + + # INPUT_SCALES + if self.quant_config.activation_scheme == "static": + if not self.quant_config.is_checkpoint_fp8_serialized: + raise ValueError( + "Found static activation scheme for checkpoint that " + "was not serialized fp8." + ) + + w13_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + + else: + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer: Module) -> None: + if _is_hip and _use_hip_int4: + self.process_weights_hip_int4(layer) + return + + # Block quant doesn't need to process weights after loading + if self.block_quant: + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_fp8_fnuz: + # activation_scheme: dynamic + w13_weight, w13_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.w13_weight, + weight_scale=layer.w13_weight_scale_inv, + input_scale=None, + ) + w2_weight, w2_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.w2_weight, + weight_scale=layer.w2_weight_scale_inv, + input_scale=None, + ) + # Reset the parameter + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w13_weight_scale_inv = torch.nn.Parameter( + w13_weight_scale, requires_grad=False + ) + layer.w13_input_scale = None + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + layer.w2_weight_scale_inv = torch.nn.Parameter( + w2_weight_scale, requires_grad=False + ) + layer.w2_input_scale = None + + if _use_aiter: + # Pre-shuffle weights + layer.w13_weight.data = shuffle_weight( + layer.w13_weight.contiguous(), (16, 16) + ) + layer.w2_weight.data = shuffle_weight( + layer.w2_weight.contiguous(), (16, 16) + ) + + if _is_cpu: + assert ( + _is_cpu_amx_available + ), "Fp8MoEMethod on CPU requires that CPU has AMX support" + _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"]) + + return + + # If checkpoint is fp16 or bfloat16, quantize in place. + if not self.quant_config.is_checkpoint_fp8_serialized: + # If ROCm, fp8_dtype will be float8_e4m3fnuz (MI300x HW) + w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype) + w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype) + + # Re-initialize w13_scale because we directly quantize + # merged w13 weights and generate a single scaling factor. + layer.w13_weight_scale = torch.nn.Parameter( + torch.ones( + layer.num_local_experts, + dtype=torch.float32, + device=w13_weight.device, + ), + requires_grad=False, + ) + for expert in range(layer.num_local_experts): + w13_weight[expert, :, :], layer.w13_weight_scale[expert] = ( + scaled_fp8_quant(layer.w13_weight.data[expert, :, :]) + ) + w2_weight[expert, :, :], layer.w2_weight_scale[expert] = ( + scaled_fp8_quant(layer.w2_weight.data[expert, :, :]) + ) + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + + if _is_hip: + self.process_weights_hip_scale_padding(layer) + return + + # If checkpoint is fp8, we need to handle that the + # MoE kernels require single activation scale and single weight + # scale for w13 per expert. + else: + # Fp8 moe kernels require a single activation scale. + # We take the max of all the scales in case they differ. + if self.quant_config.activation_scheme == "static": + if layer.w13_input_scale is None or layer.w2_input_scale is None: + raise ValueError( + "QuantConfig has static quantization, but found " + "activation scales are None." + ) + if not all_close_1d(layer.w13_input_scale) or not all_close_1d( + layer.w2_input_scale + ): + print_warning_once( + "Found input_scales that are not equal for " + "fp8 MoE layer. Using the maximum across experts " + "for each layer. " + ) + layer.w13_input_scale = torch.nn.Parameter( + layer.w13_input_scale.max(), requires_grad=False + ) + layer.w2_input_scale = torch.nn.Parameter( + layer.w2_input_scale.max(), requires_grad=False + ) + + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_fp8_fnuz: + # Normalize the weights and scales + w13_weight, w13_weight_scale, w13_input_scale = ( + normalize_e4m3fn_to_e4m3fnuz( + layer.w13_weight, layer.w13_weight_scale, layer.w13_input_scale + ) + ) + w2_weight, w2_weight_scale, w2_input_scale = ( + normalize_e4m3fn_to_e4m3fnuz( + layer.w2_weight, layer.w2_weight_scale, layer.w2_input_scale + ) + ) + # Reset the parameter + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter( + w13_weight_scale, requires_grad=False + ) + if w13_input_scale is not None: + layer.w13_input_scale = torch.nn.Parameter( + w13_input_scale, requires_grad=False + ) + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter( + w2_weight_scale, requires_grad=False + ) + if w2_input_scale is not None: + layer.w2_input_scale = torch.nn.Parameter( + w2_input_scale, requires_grad=False + ) + # Fp8 moe kernel needs single weight scale for w13 per expert. + # We take the max then dequant and requant each expert. + assert layer.w13_weight_scale is not None + shard_size = layer.intermediate_size_per_partition + max_w13_scales = layer.w13_weight_scale.max(dim=1).values + for expert_id in range(layer.num_local_experts): + start = 0 + for shard_id in range(2): + dq_weight = per_tensor_dequantize( + layer.w13_weight[expert_id][start : start + shard_size, :], + layer.w13_weight_scale[expert_id][shard_id], + ) + ( + layer.w13_weight[expert_id][start : start + shard_size, :], + _, + ) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id]) + start += shard_size + + layer.w13_weight_scale = torch.nn.Parameter( + max_w13_scales, requires_grad=False + ) + + if _is_hip: + self.process_weights_hip_scale_padding(layer) + return + + def process_weights_hip_int4(self, layer: Module): + # TODO: _use_aiter: add after triton kernel added + # INT4-FP8 (INT4 MoE Weight, FP8 Compute) + # Weight Permutation + layer.w13_weight = torch.nn.Parameter( + shuffle_weight(layer.w13_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + layer.w2_weight = torch.nn.Parameter( + shuffle_weight(layer.w2_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + + # INT4-FP8 : offset INT4 w13_weight_scale1 to single w13_weight_scale + # Fp8 moe kernel needs single fp8 w13_weight_scale for w13 per expert. + # We won't do requant each expert's fp8 weight (not direct available), + # instead we adjust half of INT4 w13_weight_scale1 numbers + assert layer.w13_weight_scale is not None + shard_size = layer.intermediate_size_per_partition + max_w13_scales = layer.w13_weight_scale.max(dim=1).values + for expert_id in range(layer.num_local_experts): + start = 0 + max_w13_scale_fp8 = max_w13_scales[expert_id] + for shard_id in range(2): + if layer.w13_weight_scale[expert_id][shard_id] != max_w13_scale_fp8: + int4_rescale = ( + layer.w13_weight_scale[expert_id][shard_id] / max_w13_scale_fp8 + ) + layer.w13_weight_scale1[expert_id][ + start : start + shard_size + ] *= int4_rescale + start += shard_size + + layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales, requires_grad=False) + + # special hack to asm_moe, which takes (weight_scale1 * weight_scale) as post GEMM scaling + # optimal design - shall apply per-column weight_scale1 before GEMM, and weight_scale post + for expert_id in range(layer.num_local_experts): + layer.w13_weight_scale1[expert_id] *= max_w13_scales[expert_id] + layer.w2_weight_scale1[expert_id] *= layer.w2_weight_scale[expert_id] + + def process_weights_hip_scale_padding(self, layer: Module): + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + padding_size, # Avoid circular import + ) + + if _use_aiter: + layer.w13_weight = torch.nn.Parameter( + shuffle_weight(layer.w13_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + layer.w2_weight = torch.nn.Parameter( + shuffle_weight(layer.w2_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + + # ROCm (_use_aiter): using column-wise scaling + layer.w13_weight_scale1 *= layer.w13_weight_scale.unsqueeze(-1) + layer.w2_weight_scale1 *= layer.w2_weight_scale.unsqueeze(-1) + elif get_bool_env_var("SGLANG_MOE_PADDING"): + # If ROCm, apply weight padding (min. Mem channel contention) only if set + layer.w13_weight = torch.nn.Parameter( + F.pad(layer.w13_weight.data, (0, padding_size), "constant", 0), + requires_grad=False, + ) + torch.cuda.empty_cache() + layer.w2_weight = torch.nn.Parameter( + F.pad(layer.w2_weight.data, (0, padding_size), "constant", 0), + requires_grad=False, + ) + torch.cuda.empty_cache() + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: DispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + moe_runner_config = self.moe_runner_config + + if use_intel_amx_backend(layer): + from sglang.srt.layers.moe.topk import apply_topk_weights_cpu + + topk_weights, topk_ids, _ = topk_output + x, topk_weights = apply_topk_weights_cpu( + moe_runner_config.apply_router_weight_on_input, topk_weights, x + ) + + output = torch.ops.sgl_kernel.fused_experts_cpu( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + False, # inplace See [Note] inplace should be False in fused_experts. + False, # use_int8_w8a8 + True, # use_fp8_w8a16 + layer.w13_weight_scale_inv, # w1_scale + layer.w2_weight_scale_inv, # w2_scale + self.quant_config.weight_block_size, # block_size + None, # a1_scale + None, # a2_scale + True, # is_vnni + ) + return StandardCombineInput(hidden_states=output) + + if _is_hip: + ret = self.maybe_apply_hip_fused_experts( + layer, + x, + topk_output, + moe_runner_config.activation, + moe_runner_config.no_combine, + ) + if ret is not None: + return StandardCombineInput(hidden_states=ret) + + if self.use_cutlass_fused_experts_fp8: + from sglang.srt.layers.moe.cutlass_moe import cutlass_fused_experts_fp8 + + topk_weights, topk_ids, _ = topk_output + output = cutlass_fused_experts_fp8( + x, + layer.w13_weight.transpose(1, 2), + layer.w2_weight.transpose(1, 2), + layer.w13_weight_scale_inv.transpose(1, 2), + layer.w2_weight_scale_inv.transpose(1, 2), + topk_weights, + topk_ids, + self.ab_strides1, + self.c_strides1, + self.ab_strides2, + self.c_strides2, + self.workspace, + self.a_ptr, + self.b_ptr, + self.out_ptr, + self.a_scales_ptr, + self.b_scales_ptr, + self.expert_offsets, + self.problem_sizes1, + self.problem_sizes2, + use_fp8_blockscale=True, + ) + return StandardCombineInput(hidden_states=output) + + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + use_fp8_w8a8=True, + w13_scale=( + layer.w13_weight_scale_inv + if self.block_quant + else layer.w13_weight_scale + ), + w2_scale=( + layer.w2_weight_scale_inv if self.block_quant else layer.w2_weight_scale + ), + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + block_shape=self.quant_config.weight_block_size, + ) + return self.runner.run(dispatch_output, quant_info) + + def apply_with_router_logits( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> torch.Tensor: + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + activation = self.moe_runner_config.activation + routed_scaling_factor = self.moe_runner_config.routed_scaling_factor + + from flashinfer.fused_moe import trtllm_fp8_block_scale_moe + + from sglang.srt.layers.moe.topk import TopKOutputChecker + + assert TopKOutputChecker.format_is_bypassed(topk_output) + router_logits = topk_output.router_logits + topk_config = topk_output.topk_config + assert ( + activation == "silu" + ), "Only silu is supported for flashinfer blockscale fp8 moe" + a_q, a_sf = per_token_group_quant_fp8(x, self.quant_config.weight_block_size[1]) + # NOTE: scales of hidden states have to be transposed! + a_sf_t = a_sf.t().contiguous() + + assert ( + topk_config.num_expert_group is not None + and topk_config.topk_group is not None + ), "Current trtllm_fp8_block_scale_moe kernel does not support these two arguments as None" + + correction_bias = ( + None + if topk_config.correction_bias is None + else topk_config.correction_bias.to(x.dtype) + ) + + return trtllm_fp8_block_scale_moe( + routing_logits=router_logits.to(torch.float32), + routing_bias=correction_bias, + hidden_states=a_q, + hidden_states_scale=a_sf_t, + gemm1_weights=layer.w13_weight, + gemm1_weights_scale=layer.w13_weight_scale_inv, + gemm2_weights=layer.w2_weight, + gemm2_weights_scale=layer.w2_weight_scale_inv, + num_experts=layer.num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=layer.w2_weight.shape[2], + local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, + local_num_experts=layer.num_local_experts, + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + tile_tokens_dim=get_tile_tokens_dim( + x.shape[0], topk_config.top_k, layer.num_experts + ), + routing_method_type=2, # DeepSeek-styled routing method + use_shuffled_weight=False, + ) + + def maybe_apply_hip_fused_experts( + self, + layer: torch.nn.Module, + x: torch.Tensor, + topk_output: TopKOutput, + activation: str = "silu", + no_combine: bool = False, + ) -> Optional[torch.Tensor]: + topk_weights, topk_ids, _ = topk_output + if _use_hip_int4: + # TODO: add triton kernel and add check _use_aiter + assert not no_combine, f"{no_combine=} is not supported." + return fused_moe( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + quant_type=QuantType.per_Token, + w1_scale=layer.w13_weight_scale1, + w2_scale=layer.w2_weight_scale1, + activation=( + ActivationType.Silu if activation == "silu" else ActivationType.Gelu + ), + ) + + if _use_aiter: + assert not no_combine, f"{no_combine=} is not supported." + if self.block_quant: + return fused_moe( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + w1_scale=layer.w13_weight_scale_inv, + w2_scale=layer.w2_weight_scale_inv, + quant_type=QuantType.per_128x128, + activation=( + ActivationType.Silu + if activation == "silu" + else ActivationType.Gelu + ), + expert_mask=None, + ) + else: + return fused_moe( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + quant_type=QuantType.per_Token, + w1_scale=layer.w13_weight_scale1, + w2_scale=layer.w2_weight_scale1, + activation=( + ActivationType.Silu + if activation == "silu" + else ActivationType.Gelu + ), + ) + return None + + +class Fp8KVCacheMethod(BaseKVCacheMethod): + """ + Supports loading kv-cache scaling factors from FP8 checkpoints. + """ + + def __init__(self, quant_config: Fp8Config): + super().__init__(quant_config) diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py new file mode 100644 index 000000000..f0512365b --- /dev/null +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -0,0 +1,1806 @@ +# Copyright 2024 SGLang 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 functools +import json +import logging +import os +from functools import lru_cache +from typing import Any, Dict, List, Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.quantization import deep_gemm_wrapper +from sglang.srt.utils import ( + align, + direct_register_custom_op, + get_bool_env_var, + get_device_core_count, + get_device_name, + is_cpu, + is_cuda, + is_hip, + log_info_on_rank0, + supports_custom_op, +) + +_is_hip = is_hip() +_is_cuda = is_cuda() +_is_cpu = is_cpu() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _is_cuda: + from sgl_kernel import ( + sgl_per_tensor_quant_fp8, + sgl_per_token_group_quant_fp8, + sgl_per_token_quant_fp8, + ) + +if _is_hip: + if _use_aiter: + try: + from aiter import ( # v0.1.3 + dynamic_per_tensor_quant, + dynamic_per_token_scaled_quant, + static_per_tensor_quant, + ) + except ImportError: + raise ImportError("aiter is required when SGLANG_USE_AITER is set to True") + else: + try: + import vllm._C + except ImportError: + raise ImportError("vllm is required when SGLANG_USE_AITER is set to False") + +logger = logging.getLogger(__name__) + + +@lru_cache() +def is_fp8_fnuz() -> bool: + if _is_hip: + # only device 0 is checked, this assumes MI300 platforms are homogeneous + return "gfx94" in torch.cuda.get_device_properties(0).gcnArchName + return False + + +if is_fp8_fnuz(): + fp8_dtype = torch.float8_e4m3fnuz + fp8_max = 224.0 +else: + fp8_dtype = torch.float8_e4m3fn + fp8_max = torch.finfo(fp8_dtype).max +fp8_min = -fp8_max + +if supports_custom_op(): + + def deep_gemm_fp8_fp8_bf16_nt( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, + ) -> None: + deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C) + + def deep_gemm_fp8_fp8_bf16_nt_fake( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, + ) -> None: + return + + direct_register_custom_op( + op_name="deep_gemm_fp8_fp8_bf16_nt", + op_func=deep_gemm_fp8_fp8_bf16_nt, + mutates_args=["C"], + fake_impl=deep_gemm_fp8_fp8_bf16_nt_fake, + ) + + +@triton.jit +def _per_token_group_quant_8bit( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + # Stride of input + y_stride, + # Columns of input + N, + # Avoid to divide zero + eps, + # Information for float8 + bit8_min, + bit8_max, + # Meta-parameters + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group quantization on a + tensor. + + This function converts the tensor values into float8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * y_stride + y_q_ptr += g_id * y_stride + y_s_ptr += g_id + + cols = tl.arange(0, BLOCK) # N <= BLOCK + mask = cols < N + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / bit8_max + y_s_inv = 1.0 / y_s + y_q = tl.clamp(y * y_s_inv, bit8_min, bit8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + + +@triton.jit +def _per_token_group_quant_8bit_colmajor( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + group_size, + # Num columns of y + y_num_columns, + # Stride from one column to the next of y_s + y_s_col_stride, + # Avoid to divide zero + eps, + # Information for float8 + bit8_min, + bit8_max, + # Meta-parameters + BLOCK: tl.constexpr, + SCALE_UE8M0: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group + quantization on a tensor. + This function converts the tensor values into float8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id.to(tl.int64) * group_size + y_q_ptr += g_id.to(tl.int64) * group_size + + # Convert g_id the flattened block coordinate to 2D so we can index + # into the output y_scales matrix + blocks_per_row = y_num_columns // group_size + scale_col = g_id % blocks_per_row + scale_row = g_id // blocks_per_row + y_s_ptr += scale_col * y_s_col_stride + scale_row + + cols = tl.arange(0, BLOCK) # group_size <= BLOCK + mask = cols < group_size + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / bit8_max + if SCALE_UE8M0: + y_s = tl.exp2(tl.ceil(tl.log2(tl.abs(y_s)))) + y_q = tl.clamp(y / y_s, bit8_min, bit8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + + +def _per_token_group_quant_8bit_raw( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: torch.dtype = fp8_dtype, + column_major_scales: bool = False, + scale_tma_aligned: bool = False, + scale_ue8m0: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Function to perform per-token-group quantization on an input tensor `x`. + + It converts the tensor values into signed float8 values and returns the + quantized tensor along with the scaling factor used for quantization. + + Args: + x: The input tenosr with ndim >= 2. + group_size: The group size used for quantization. + eps: The minimum to avoid dividing zero. + dtype: The dype of output tensor. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. + """ + assert ( + x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.is_contiguous(), "`x` is not contiguous" + + if _is_hip: + if dtype == torch.int8: + bit8_max = 127.0 + else: + bit8_max = 224.0 + bit8_min = -bit8_max # TODO incorrect for int8 + else: + if dtype == torch.int8: + info = torch.iinfo(dtype) + else: + info = torch.finfo(dtype) + bit8_max = info.max + bit8_min = info.min + + x_q = torch.empty_like(x, device=x.device, dtype=dtype) + x_s = create_per_token_group_quant_fp8_output_scale( + x_shape=x.shape, + device=x.device, + group_size=group_size, + column_major_scales=column_major_scales, + scale_tma_aligned=scale_tma_aligned, + scale_ue8m0=False, + ) + + M = x.numel() // group_size + N = group_size + + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + if column_major_scales: + _per_token_group_quant_8bit_colmajor[(M,)]( + x, + x_q, + x_s, + group_size, + x.shape[1], + x_s.stride(1), + eps, + bit8_min=bit8_min, + bit8_max=bit8_max, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=num_stages, + SCALE_UE8M0=scale_ue8m0, + ) + else: + assert not scale_ue8m0 + _per_token_group_quant_8bit[(M,)]( + x, + x_q, + x_s, + group_size, + N, + eps, + bit8_min=bit8_min, + bit8_max=bit8_max, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=num_stages, + ) + + if scale_ue8m0: + from deep_gemm import transform_sf_into_required_layout + + assert group_size == 128 + x_s = transform_sf_into_required_layout( + x_s, + num_groups=None, + mn=x_q.shape[0], + k=x_q.shape[1], + recipe=(1, group_size, group_size), + is_sfa=True, + ) + + return x_q, x_s + + +# backward compatibility +per_token_group_quant_fp8 = _per_token_group_quant_8bit_raw + + +def _per_token_group_quant_8bit_fuse_silu_and_mul( + x: torch.Tensor, + group_size: int, + dst_dtype: torch.dtype, + column_major_scales: bool, + scale_tma_aligned: bool, + scale_ue8m0: bool, + masked_m: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + # Another way to implement (can be used in e.g. comparison tests) + # from sgl_kernel import silu_and_mul + # x_after_silu_and_mul = silu_and_mul(x) + # return per_token_group_quant_fp8( + # x_after_silu_and_mul, + # group_size=group_size, + # eps=eps, + # column_major_scales=column_major_scales, + # scale_tma_aligned=scale_tma_aligned, + # scale_ue8m0=scale_ue8m0, + # ) + + from deep_gemm import transform_sf_into_required_layout + + from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_masked_post_quant_fwd + + assert column_major_scales + assert scale_tma_aligned + assert scale_ue8m0 + + needs_unsqueeze = x.dim() == 2 + if needs_unsqueeze: + num_tokens, _ = x.shape + x = x.unsqueeze(0) + assert masked_m is None + masked_m = torch.tensor([num_tokens], device=x.device, dtype=torch.int32) + + # Use `zeros` for easier testing + output = torch.zeros( + (*x.shape[:-1], x.shape[-1] // 2), + device=x.device, + dtype=dst_dtype, + ) + # Use `zeros` for easier testing + output_scale_for_kernel = torch.zeros( + (*x.shape[:-1], x.shape[-1] // 2 // group_size), + device=x.device, + dtype=torch.float32, + ) + silu_and_mul_masked_post_quant_fwd( + input=x, + output=output, + output_scale=output_scale_for_kernel, + quant_group_size=group_size, + masked_m=masked_m, + scale_ue8m0=scale_ue8m0, + ) + + assert group_size == 128 + output_scale = transform_sf_into_required_layout( + output_scale_for_kernel, + num_groups=output.shape[0], + mn=output.shape[-2], + k=output.shape[-1], + recipe=(1, group_size, group_size), + is_sfa=True, + ) + + if needs_unsqueeze: + output = output.squeeze(0) + output_scale = output_scale.squeeze(0) + + return output, output_scale + + +def per_token_group_quant_8bit( + x: torch.Tensor, + group_size: int, + dst_dtype: torch.dtype, + eps: float = 1e-10, + column_major_scales: bool = False, + scale_tma_aligned: bool = False, + scale_ue8m0: bool = False, + fuse_silu_and_mul: bool = False, + masked_m: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + if fuse_silu_and_mul: + return _per_token_group_quant_8bit_fuse_silu_and_mul( + x=x, + group_size=group_size, + dst_dtype=dst_dtype, + column_major_scales=column_major_scales, + scale_tma_aligned=scale_tma_aligned, + scale_ue8m0=scale_ue8m0, + masked_m=masked_m, + ) + else: + return _per_token_group_quant_8bit_raw( + x=x, + group_size=group_size, + eps=eps, + column_major_scales=column_major_scales, + scale_tma_aligned=scale_tma_aligned, + scale_ue8m0=scale_ue8m0, + dtype=dst_dtype, + ) + + +def create_per_token_group_quant_fp8_output_scale( + x_shape, + device, + group_size, + column_major_scales: bool, + scale_tma_aligned: bool, + scale_ue8m0: bool, +): + if scale_ue8m0: + assert column_major_scales and scale_tma_aligned + *x_batch, x_q_mn, x_q_k = x_shape + x_s_mn, x_s_k = x_q_mn, x_q_k // 128 + aligned_mn = align(x_s_mn, 4) + aligned_k = align(x_s_k, 4) + # TODO(FIXME): Fix cuda kernel and recover here to empty. + return torch.empty( + (*x_batch, aligned_k // 4, aligned_mn), + device=device, + dtype=torch.int, + ).transpose(-1, -2)[..., :x_s_mn, :] + elif column_major_scales: + if scale_tma_aligned: + # TODO extract "align" function + # aligned to 4 * sizeof(float) + aligned_size = (x_shape[-2] + 3) // 4 * 4 + return torch.empty( + x_shape[:-2] + (x_shape[-1] // group_size, aligned_size), + device=device, + dtype=torch.float32, + ).permute(-1, -2)[: x_shape[-2], :] + else: + return torch.empty( + (x_shape[-1] // group_size,) + x_shape[:-1], + device=device, + dtype=torch.float32, + ).permute(-1, -2) + else: + return torch.empty( + x_shape[:-1] + (x_shape[-1] // group_size,), + device=device, + dtype=torch.float32, + ) + + +def sglang_per_token_group_quant_fp8( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + column_major_scales: bool = False, + scale_tma_aligned: bool = False, + scale_ue8m0: bool = False, + fuse_silu_and_mul: bool = False, + masked_m: Optional[torch.Tensor] = None, +): + assert ( + x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.is_contiguous(), "`x` is not contiguous" + + out_shape = (*x.shape[:-1], x.shape[-1] // (2 if fuse_silu_and_mul else 1)) + + x_q = torch.empty(out_shape, device=x.device, dtype=fp8_dtype) + x_s = create_per_token_group_quant_fp8_output_scale( + x_shape=out_shape, + device=x.device, + group_size=group_size, + column_major_scales=column_major_scales, + scale_tma_aligned=scale_tma_aligned, + scale_ue8m0=scale_ue8m0, + ) + + if x.shape[0] > 0: + sgl_per_token_group_quant_fp8( + x, x_q, x_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0 + ) + + return x_q, x_s + + +# TODO maybe unify int8 and fp8 code later +def sglang_per_token_group_quant_8bit( + x: torch.Tensor, + group_size: int, + dst_dtype: torch.dtype, + eps: float = 1e-10, + column_major_scales: bool = False, + scale_tma_aligned: bool = False, + scale_ue8m0: bool = False, + fuse_silu_and_mul: bool = False, + masked_m: Optional[torch.Tensor] = None, +): + from sglang.srt.layers.quantization.int8_kernel import ( + sglang_per_token_group_quant_int8, + ) + + if dst_dtype == torch.int8: + assert not column_major_scales + assert not scale_tma_aligned + assert not fuse_silu_and_mul + assert masked_m is None + return sglang_per_token_group_quant_int8( + x=x, + group_size=group_size, + eps=eps, + dtype=dst_dtype, + ) + + return sglang_per_token_group_quant_fp8( + x=x, + group_size=group_size, + eps=eps, + column_major_scales=column_major_scales, + scale_tma_aligned=scale_tma_aligned, + scale_ue8m0=scale_ue8m0, + fuse_silu_and_mul=fuse_silu_and_mul, + masked_m=masked_m, + ) + + +def sglang_per_token_quant_fp8( + x: torch.Tensor, + dtype: torch.dtype = fp8_dtype, +): + assert x.is_contiguous(), "`x` is not contiguous" + + x_q = torch.empty_like(x, device=x.device, dtype=dtype) + x_s = torch.empty( + x.shape[0], + 1, + device=x.device, + dtype=torch.float32, + ) + + sgl_per_token_quant_fp8(x, x_q, x_s) + + return x_q, x_s + + +@triton.jit +def _static_quant_fp8( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + y_s_repeat_ptr, + # Stride of input + y_stride, + # Columns of input + N, + # Information for float8 + fp8_min, + fp8_max, + # Meta-parameters + BLOCK: tl.constexpr, + REPEAT_SCALE: tl.constexpr, +): + """A Triton-accelerated function to perform quantization using the given scale on a + tensor + + This function converts the tensor values into float8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * y_stride + y_q_ptr += g_id * y_stride + if REPEAT_SCALE: + y_s_repeat_ptr += g_id + + cols = tl.arange(0, BLOCK) # N <= BLOCK + mask = cols < N + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + y_s = tl.load(y_s_ptr).to(tl.float32) + y_s_inv = 1.0 / y_s + y_q = tl.clamp(y * y_s_inv, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + if REPEAT_SCALE: + tl.store(y_s_repeat_ptr, y_s) + + +def static_quant_fp8( + x: torch.Tensor, + x_s: torch.Tensor, + repeat_scale: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Function to perform static quantization using the given scale on an input tensor `x`. + + It converts the tensor values into signed float8 values and returns the + quantized tensor along with the scaling factor used for quantization. + + Args: + x: The input tenosr with ndim >= 2. + x_s: The quantization scale. + repeat_scale: Whether to broadcast per-tensor scale to per-channel scale. + dtype: The dype of output tensor. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. + """ + assert x.is_contiguous(), "`x` is not contiguous" + assert x_s.numel() == 1, "only supports per-tensor scale" + + x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) + M = x.numel() // x.shape[-1] + N = x.shape[-1] + if repeat_scale: + x_s_repeat = torch.empty( + (M, 1), + device=x.device, + dtype=torch.float32, + ) + else: + x_s_repeat = None + + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + _static_quant_fp8[(M,)]( + x, + x_q, + x_s, + x_s_repeat, + N, + N, + fp8_min=fp8_min, + fp8_max=fp8_max, + BLOCK=BLOCK, + REPEAT_SCALE=repeat_scale, + num_warps=num_warps, + num_stages=num_stages, + ) + x_s = x_s_repeat if repeat_scale else x_s + return x_q, x_s + + +@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) + + +@triton.jit +def _w8a8_block_fp8_matmul_unrolledx4( + # 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) + # manually unroll to 4 iterations + UNROLL_FACTOR = 4 + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * UNROLL_FACTOR)): + # 1st iteration + a = tl.load( + a_ptrs, + mask=offs_k[None, :] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K, + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K, + other=0.0, + ) + + k_start = (k * UNROLL_FACTOR) * 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 + + # 2nd iteration + a = tl.load( + a_ptrs, + mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K, + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K, + other=0.0, + ) + + k_start = k_start + 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 + + # 3rd iteration + a = tl.load( + a_ptrs, + mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K, + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K, + other=0.0, + ) + + k_start = k_start + 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 + + # 4th iteration + a = tl.load( + a_ptrs, + mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K, + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K, + other=0.0, + ) + + k_start = k_start + 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) + + +@functools.lru_cache +def get_w8a8_block_fp8_configs( + N: int, K: int, block_n: int, block_k: int +) -> Optional[Dict[int, Any]]: + """ + Return optimized configurations for the w8a8 block fp8 kernel. + + The return value will be a dictionary that maps an irregular grid of + batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the + kernel on a given batch size bs, the closest batch size in the grid should + be picked and the associated configuration chosen to invoke the kernel. + """ + + # First look up if an optimized configuration is available in the configs + # directory + device_name = get_device_name().replace(" ", "_") + json_file_name = f"N={N},K={K},device_name={device_name},dtype=fp8_w8a8,block_shape=[{block_n}, {block_k}].json" + + config_file_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name + ) + if os.path.exists(config_file_path): + with open(config_file_path) as f: + log_info_on_rank0( + logger, + f"Using configuration from {config_file_path} for W8A8 Block FP8 kernel.", + ) + # If a configuration has been found, return it + return {int(key): val for key, val in json.load(f).items()} + + # If no optimized configuration is available, we will use the default + # configuration + logger.warning( + ( + "Using default W8A8 Block FP8 kernel config. Performance might be sub-optimal! " + "Config file not found at %s" + ), + config_file_path, + ) + return None + + +def select_w8a8_block_fp8_matmul_kernel(M, N, META): + return _w8a8_block_fp8_matmul + + +if _is_hip: + + def use_w8a8_block_fp8_matmul_unrolledx4(M, N, META): + # Use manually unrolledx4 kernel on AMD GPU when the grid size is small. + # Empirical testing shows the sweet spot lies when it's less than the # of + # compute units available on the device. + num_workgroups = triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv( + N, META["BLOCK_SIZE_N"] + ) + num_workgroups <= get_device_core_count() + + def select_w8a8_block_fp8_matmul_kernel(M, N, META): + if use_w8a8_block_fp8_matmul_unrolledx4(M, N, META): + return _w8a8_block_fp8_matmul_unrolledx4 + else: + return _w8a8_block_fp8_matmul + + +def prepare_block_fp8_matmul_inputs( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype = torch.float16, +) -> Tuple[int, int, int]: + assert len(block_size) == 2 + block_n, block_k = block_size[0], block_size[1] + + assert A.shape[-1] == B.shape[-1] + assert A.shape[:-1] == As.shape[:-1] + assert A.is_contiguous() + + if As.dtype == torch.float: + assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] + elif As.dtype == torch.int: + assert ( + triton.cdiv(triton.cdiv(A.shape[-1], block_k), 4) == As.shape[-1] + ), f"{A.shape=} {As.shape=} {block_size=}" + else: + raise NotImplementedError + + M = A.numel() // A.shape[-1] + + assert B.ndim == 2 + assert B.is_contiguous() + assert Bs.ndim == 2 + N, K = B.shape + + if Bs.dtype == torch.float: + assert triton.cdiv(N, block_n) == Bs.shape[0] + assert triton.cdiv(K, block_k) == Bs.shape[1] + elif Bs.dtype == torch.int: + assert N == Bs.shape[0], f"{B.shape=} {Bs.shape=} {block_size=}" + assert ( + triton.cdiv(triton.cdiv(K, block_k), 4) == Bs.shape[1] + ), f"{B.shape=} {Bs.shape=} {block_size=}" + else: + raise NotImplementedError + + C_shape = A.shape[:-1] + (N,) + C = A.new_empty(C_shape, dtype=output_dtype) + + return M, N, K, C + + +def w8a8_block_fp8_matmul_deepgemm( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype, +) -> torch.Tensor: + M, N, K, C = prepare_block_fp8_matmul_inputs(A, B, As, Bs, block_size, output_dtype) + + # Deepgemm only supports output tensor type as bfloat16 + assert C.dtype == torch.bfloat16 and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + + if supports_custom_op(): + torch.ops.sglang.deep_gemm_fp8_fp8_bf16_nt(A, As, B, Bs, C) + else: + deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C) + + return C + + +def w8a8_block_fp8_matmul_triton( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + """This function performs matrix multiplication with block-wise quantization. + + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + + Returns: + torch.Tensor: The result of matmul. + """ + + M, N, K, C = prepare_block_fp8_matmul_inputs(A, B, As, Bs, block_size, output_dtype) + + block_n, block_k = block_size + + configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1]) + if configs: + # If an optimal configuration map has been found, look up the + # optimal config + config = configs[min(configs.keys(), key=lambda x: abs(x - M))] + else: + # Default config + # Block-wise quant: BLOCK_SIZE_K must be divisible by block_size[1] + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": block_size[0], + "BLOCK_SIZE_K": block_size[1], + "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"]), + ) + + kernel = select_w8a8_block_fp8_matmul_kernel(M, N, config) + + kernel[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + **config, + ) + + return C + + +# universal entry point, for testing purposes +def w8a8_block_fp8_matmul( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + if output_dtype == torch.bfloat16 and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + return w8a8_block_fp8_matmul_deepgemm( + A, B, As, Bs, block_size, output_dtype=output_dtype + ) + + return w8a8_block_fp8_matmul_triton( + A, B, As, Bs, block_size, output_dtype=output_dtype + ) + + +@triton.jit +def _per_tensor_quant_mla_fp8_stage1( + x_ptr, + x_s_ptr, + head_size, + x_stride_h, + x_stride_s, + eps, + fp8_max, + BLOCK_SIZE: tl.constexpr, +): + seq_id = tl.program_id(0) + head_id = tl.program_id(1) + offset = tl.arange(0, BLOCK_SIZE) + mask = offset < head_size + + x_ptr += head_id * x_stride_h + seq_id * x_stride_s + x = tl.load(x_ptr + offset, mask=mask, other=0.0).to(tl.float32) + _absmax = tl.maximum(tl.max(tl.abs(x)), eps) + + tl.atomic_max(x_s_ptr, _absmax / fp8_max) + + +@triton.jit +def _per_tensor_quant_mla_fp8_stage2( + x_ptr, + x_s_ptr, + x_q_ptr, + num_seq, + head_size, + x_stride_h, + x_stride_s, + fp8_min, + fp8_max, + BLOCK_SIZE: tl.constexpr, +): + seq_id = tl.program_id(0) + head_id = tl.program_id(1) + offset = tl.arange(0, BLOCK_SIZE) + mask = offset < head_size + + x_s = tl.load(x_s_ptr) + x_s_inv = 1.0 / x_s + + x_ptr += head_id * x_stride_h + seq_id * x_stride_s + x_q_ptr += head_id * num_seq * head_size + seq_id * head_size + + x = tl.load(x_ptr + offset, mask=mask, other=0.0).to(tl.float32) + x_q = tl.clamp(x * x_s_inv, fp8_min, fp8_max).to(x_q_ptr.dtype.element_ty) + tl.store(x_q_ptr + offset, x_q, mask=mask) + + +def per_tensor_quant_mla_fp8( + x: torch.Tensor, x_s_out: torch.Tensor, eps: float = 1e-12 +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + This function quantizes input values to float8 values with tensor-wise quantization + and specialized for mla absorbed case. + """ + assert x.dim() == 3, "`x` is not a 3d-tensor" + assert ( + x_s_out.shape == (1,) + and x_s_out.dtype == torch.float32 + and x_s_out.device == x.device + ) + + x_q = x.new_empty(x.size(), dtype=fp8_dtype) + + num_head, num_seq, head_size = x.shape + BLOCK_SIZE = triton.next_power_of_2(head_size) + grid = (num_seq, num_head) + + _per_tensor_quant_mla_fp8_stage1[grid]( + x, + x_s_out, + head_size, + x.stride(0), + x.stride(1), + eps, + fp8_max, + BLOCK_SIZE, + ) + _per_tensor_quant_mla_fp8_stage2[grid]( + x, + x_s_out, + x_q, + num_seq, + head_size, + x.stride(0), + x.stride(1), + fp8_min, + fp8_max, + BLOCK_SIZE, + ) + + return x_q, x_s_out + + +@triton.jit +def _per_token_group_quant_mla_deep_gemm_masked_fp8( + y_ptr, + y_q_ptr, + y_s_ptr, + masked_m_ptr, + group_size, + y_stride_b, + y_stride_t, + y_q_stride_b, + y_q_stride_t, + y_s_stride_b, + y_s_stride_g, + eps, + fp8_min, + fp8_max, + NUM_GROUP: tl.constexpr, + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group + quantization on a tensor for deep_gemm grouped_gemm_masked. + This function converts the tensor values into float8 values. + y and y_q: (b, t, k) + y_s: (b, k//group_size, t) + """ + t_id = tl.program_id(0) + b_id = tl.program_id(1) + + y_ptr += b_id * y_stride_b + t_id * y_stride_t + y_q_ptr += b_id * y_q_stride_b + t_id * y_q_stride_t + y_s_ptr += b_id * y_s_stride_b + t_id + + if t_id == 0: + tl.store(masked_m_ptr + b_id, tl.num_programs(0)) + + cols = tl.arange(0, BLOCK) # group_size <= BLOCK + mask = cols < group_size + + for gid in range(NUM_GROUP): + y = tl.load(y_ptr + gid * group_size + cols, mask=mask, other=0.0).to( + tl.float32 + ) + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / fp8_max + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + gid * group_size + cols, y_q, mask=mask) + tl.store(y_s_ptr + gid * y_s_stride_g, y_s) + + +def per_token_group_quant_mla_deep_gemm_masked_fp8( + x: torch.Tensor, + group_size: int = 128, + eps: float = 1e-12, + dtype: torch.dtype = fp8_dtype, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + This function quantizes input values to float8 values with per-token-group-quantization + for deep_gemm grouped_gemm_masked and specialized for mla absorbed case. + """ + assert x.dim() == 3, "`x` is not a 3d-tensor" + + b, m, k = x.shape + aligned_m = (m + 255) // 256 * 256 # 256 is the max block_m of the gemm kernel + num_tiles_k = k // group_size + assert num_tiles_k * group_size == k, f"k % {group_size} must be zero" + + x_q = x.new_empty((b, aligned_m, k), dtype=dtype) + x_s = x.new_empty((b, num_tiles_k, aligned_m), dtype=torch.float32) + masked_m = x.new_empty((b,), dtype=torch.int32) + + BLOCK_SIZE = triton.next_power_of_2(group_size) + grid = (m, b) + + _per_token_group_quant_mla_deep_gemm_masked_fp8[grid]( + x, + x_q, + x_s, + masked_m, + group_size, + x.stride(0), + x.stride(1), + x_q.stride(0), + x_q.stride(1), + x_s.stride(0), + x_s.stride(1), + eps, + -fp8_max, + fp8_max, + num_tiles_k, + BLOCK_SIZE, + ) + + return x_q, x_s.transpose(1, 2), masked_m, m, aligned_m + + +""" +Quantize input tensor to FP8 (8-bit floating point) format. + +Args: + input (torch.Tensor): Input tensor to be quantized + scale (Optional[torch.Tensor]): Pre-computed scaling factor for static quantization. + If None, scales will be computed dynamically. + num_token_padding (Optional[int]): If specified, pad the first dimension + of the output to at least this value. + use_per_token_if_dynamic (bool): When using dynamic scaling (scale=None), + determines the quantization granularity: + - True: compute scale per token + - False: compute single scale per tensor + +Returns: + Tuple[torch.Tensor, torch.Tensor]: A tuple containing: + - quantized_tensor: The FP8 quantized version of input + - scale_tensor: The scaling factors used for quantization + +Raises: + AssertionError: If input is not 2D or if static scale's numel != 1 +""" +if _is_hip: + + def scaled_fp8_quant( + input: torch.Tensor, + scale: Optional[torch.Tensor] = None, + num_token_padding: Optional[int] = None, + use_per_token_if_dynamic: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert input.ndim == 2, f"Expected 2D input tensor, got {input.ndim}D" + shape = input.shape + if num_token_padding: + shape = (max(num_token_padding, input.shape[0]), shape[1]) + output = torch.empty(shape, device=input.device, dtype=fp8_dtype) + + if scale is None: + # Dynamic scaling + if use_per_token_if_dynamic: + scale = torch.empty( + (shape[0], 1), device=input.device, dtype=torch.float32 + ) + if _use_aiter: + dynamic_per_token_scaled_quant(output, input, scale) + else: + torch.ops._C.dynamic_per_token_scaled_fp8_quant( + output, input.contiguous(), scale, None + ) + else: + scale = torch.zeros(1, device=input.device, dtype=torch.float32) + if _use_aiter: + dynamic_per_tensor_quant(output, input, scale) + else: + torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale) + else: + # Static scaling + assert ( + scale.numel() == 1 + ), f"Expected scalar scale, got numel={scale.numel()}" + if _use_aiter: + static_per_tensor_quant(output, input, scale) + else: + torch.ops._C.static_scaled_fp8_quant(output, input, scale) + + return output, scale + +else: + + def scaled_fp8_quant( + input: torch.Tensor, + scale: Optional[torch.Tensor] = None, + num_token_padding: Optional[int] = None, + use_per_token_if_dynamic: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + + assert input.ndim == 2, f"Expected 2D input tensor, got {input.ndim}D" + shape = input.shape + if num_token_padding: + shape = (max(num_token_padding, input.shape[0]), shape[1]) + output = torch.empty(shape, device=input.device, dtype=fp8_dtype) + + if scale is None: + # Dynamic scaling + if use_per_token_if_dynamic: + scale = torch.empty( + (shape[0], 1), device=input.device, dtype=torch.float32 + ) + sgl_per_token_quant_fp8(input, output, scale) + else: + scale = torch.zeros(1, device=input.device, dtype=torch.float32) + sgl_per_tensor_quant_fp8( + input, output, scale, is_static=False + ) # False for dynamic + else: + # Static scaling + assert ( + scale.numel() == 1 + ), f"Expected scalar scale, got numel={scale.numel()}" + sgl_per_tensor_quant_fp8( + input, output, scale, is_static=True + ) # True for static + + return output, scale + + +fp8_autotune = triton.autotune( + configs=[ + triton.Config({"BLOCK_M": block_m}, num_warps=num_warps) + for block_m in [16, 32, 64, 128] + for num_warps in [2, 4, 8] + ], + key=["K", "BLOCK_K", "M_ALIGNMENT"], +) + + +@triton.jit +def _per_token_group_quant_fp8_hopper_moe_mn_major( + a, # (M, K):(K, 1) + expert_offsets, # (num_experts,) + problem_sizes, # (num_experts, 3) + a_fp8, # (M, K):(K, 1) + sfa, # (M, k) + K: tl.constexpr, + BLOCK_K: tl.constexpr, + M_ALIGNMENT: tl.constexpr, + BLOCK_M: tl.constexpr, # tune +): + k_offset = tl.program_id(0) + expert_id = tl.program_id(1) + + m = tl.load(problem_sizes + expert_id * 3) + current_expert_offset = tl.load(expert_offsets + expert_id).to(tl.int64) + tl.multiple_of(m, M_ALIGNMENT) + tl.multiple_of(current_expert_offset, M_ALIGNMENT) + + coord_k = k_offset * BLOCK_K + tl.arange(0, BLOCK_K) + for i in tl.range(tl.cdiv(m, BLOCK_M)): + coord_m = i * BLOCK_M + tl.arange(0, BLOCK_M) + a_ptrs = a + current_expert_offset * K + coord_m[:, None] * K + coord_k[None, :] + a_mask = (coord_m < m)[:, None] & (coord_k < K)[None, :] + + inp = tl.load(a_ptrs, mask=a_mask).to(tl.float32) # [BLOCK_M, BLOCK_K] + inp_amax = tl.max(tl.abs(inp), axis=1) # [BLOCK_M,] + inp_amax = tl.clamp(inp_amax, min=1e-4, max=float("inf")) + inp_fp8 = (inp * (448.0 / inp_amax[:, None])).to(tl.float8e4nv) + + # Store fp8 + a_fp8_ptrs = ( + a_fp8 + current_expert_offset * K + coord_m[:, None] * K + coord_k[None, :] + ) + tl.store(a_fp8_ptrs, inp_fp8, mask=a_mask) + + # Store sfa + k = tl.cdiv(K, BLOCK_K) + sfa_ptrs = ( + sfa + current_expert_offset * k + k_offset * m + coord_m + ) # MN-Major with sfa + tl.store(sfa_ptrs, inp_amax / 448.0, mask=coord_m < m) + + +if not _is_cpu: + _per_token_group_quant_fp8_hopper_moe_mn_major = fp8_autotune( + _per_token_group_quant_fp8_hopper_moe_mn_major + ) + + +def per_token_group_quant_fp8_hopper_moe_mn_major( + A: torch.Tensor, + expert_offsets: torch.Tensor, + problem_sizes: torch.Tensor, + group_size: int, + expert_tokens_alignment: int = 1, +) -> Tuple[torch.Tensor, torch.Tensor]: + assert A.dim() == 2 + assert A.is_contiguous(), "`A` is not contiguous" + assert ( + A.shape[-1] % group_size == 0 + ), "the last dimension of `A` cannot be divisible by `group_size`" + + a_q = torch.empty_like(A, device=A.device, dtype=fp8_dtype) + M, K = A.shape[0], A.shape[1] + k = K // group_size + sfa = torch.empty((M, k), device=A.device, dtype=torch.float32) + num_experts = problem_sizes.shape[0] + grid = (k, num_experts) + _per_token_group_quant_fp8_hopper_moe_mn_major[grid]( + A, + expert_offsets, + problem_sizes, + a_q, + sfa, + K, + group_size, + expert_tokens_alignment, + ) + return a_q, sfa + + +@triton.jit +def _per_group_transpose( + data_ptr: torch.Tensor, + trans_data_ptr: torch.Tensor, + expert_offsets: torch.Tensor, + k: int, + M_ALIGNMENT: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + expert_id = tl.program_id(0) + m_id = tl.program_id(1) + k_id = tl.program_id(2) + + curr_expert_offset = tl.load(expert_offsets + expert_id) + next_expert_offset = tl.load(expert_offsets + expert_id + 1) + num_tokens_of_expert = next_expert_offset - curr_expert_offset + tl.multiple_of(curr_expert_offset, M_ALIGNMENT) + tl.multiple_of(next_expert_offset, M_ALIGNMENT) + + data_start_ptr = data_ptr + curr_expert_offset * k + trans_data_start_ptr = trans_data_ptr + curr_expert_offset * k + + k_coord = k_id * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + k_mask = k_coord < k + for start_m in tl.range(0, num_tokens_of_expert, BLOCK_SIZE_M * tl.num_programs(1)): + m_coord = start_m + m_id * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + m_mask = m_coord < num_tokens_of_expert + off = m_coord[:, None] * k + k_coord[None, :] + trans_off = m_coord[:, None] + k_coord[None, :] * num_tokens_of_expert + mask = m_mask[:, None] & k_mask[None, :] + + data = tl.load(data_start_ptr + off, mask=mask) + tl.store(trans_data_start_ptr + trans_off, data, mask=mask) + + +def per_group_transpose( + a: torch.Tensor, + expert_offsets: torch.Tensor, + M_ALIGNMENT: int = 1, +) -> torch.Tensor: + assert a.dim() == 2 + assert a.is_contiguous(), "`a` is not contiguous" + + m, k = a.size() + trans_a = torch.empty_like(a) + num_experts = expert_offsets.size(0) - 1 + + grid = lambda META: ( + num_experts, + triton.cdiv((m + num_experts - 1) // num_experts, META["BLOCK_SIZE_M"]), + triton.cdiv(k, META["BLOCK_SIZE_K"]), + ) + _per_group_transpose[grid]( + a, trans_a, expert_offsets, k, M_ALIGNMENT, BLOCK_SIZE_M=16, BLOCK_SIZE_K=8 + ) + return trans_a + + +def is_weak_contiguous(x: torch.Tensor): + strides = x.stride() + sizes = x.shape + is_not_transpose = strides[0] == 1 and (strides[1] >= max(1, sizes[0])) + is_transpose = strides[1] == 1 and (strides[0] >= max(1, sizes[1])) + return is_transpose or is_not_transpose + + +@triton.jit +def scaled_mm_kernel( + a_ptr, + b_ptr, + scale_a_ptr, + scale_b_ptr, + c_ptr, + bias_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + ACCUMULATOR_DTYPE: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_SCALE_A: tl.constexpr, + BLOCK_SIZE_SCALE_B: tl.constexpr, +): + pid = tl.program_id(axis=0) + + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + accumulator_dtype = ACCUMULATOR_DTYPE + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype) + + # NOTE: Some tensor inputs are so large, they will cause int32 overflow + # so it is necessary to use tl.int64 for all the offsets, else SEGV will + # eventually occur. + + # Offsets and masks. + offsets_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + masks_am = offsets_am < M + + offsets_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64) + masks_bn = offsets_bn < N + + offsets_k = tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + offsets_a = stride_am * offsets_am[:, None] + stride_ak * offsets_k[None, :] + offsets_b = stride_bk * offsets_k[:, None] + stride_bn * offsets_bn[None, :] + + # NOTE: BLOCK_SIZE_SCALE_A could be 1 or BLOCK_SIZE_M, so need to create + # appropriate offsets and masks for each case. Same goes for + # BLOCK_SIZE_SCALE_B. + offsets_scale_am = ( + tl.arange(0, BLOCK_SIZE_SCALE_A) + + (BLOCK_SIZE_SCALE_A > 1) * pid_m * BLOCK_SIZE_M + ) + masks_scale_am = offsets_scale_am < M + + offsets_scale_bn = ( + tl.arange(0, BLOCK_SIZE_SCALE_B) + + (BLOCK_SIZE_SCALE_B > 1) * pid_n * BLOCK_SIZE_N + ) + masks_scale_bn = offsets_scale_bn < N + + a_ptrs = a_ptr + offsets_a + b_ptrs = b_ptr + offsets_b + + scale_a_ptrs = scale_a_ptr + offsets_scale_am + scale_b_ptrs = scale_b_ptr + offsets_scale_bn + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + masks_k = offsets_k < K + masks_a = masks_am[:, None] & masks_k[None, :] + a = tl.load(a_ptrs, mask=masks_a) + + masks_b = masks_k[:, None] & masks_bn[None, :] + b = tl.load(b_ptrs, mask=masks_b) + + # Accumulate results. + accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype) + + offsets_k += BLOCK_SIZE_K + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # Apply scale at end. + masks_scale_a = masks_scale_am[:, None] & (tl.arange(0, 1) < 1)[:, None] + scale_a = tl.load(scale_a_ptrs[:, None], masks_scale_a) + # Need to broadcast to the appropriate size, if scale_a is already + # (BLOCK_SIZE_M, 1) then it will broadcast to its own shape. Same goes + # for scale_b below. + scale_a = scale_a.broadcast_to((BLOCK_SIZE_M, 1)) + accumulator = scale_a * accumulator.to(tl.float32) + + masks_scale_b = masks_scale_bn[:, None] & (tl.arange(0, 1) < 1)[None, :] + scale_b = tl.load(scale_b_ptrs[:, None], masks_scale_b) + scale_b = scale_b.broadcast_to((BLOCK_SIZE_N, 1)) + accumulator = scale_b.T * accumulator.to(tl.float32) + + # Convert to output format. + c = accumulator.to(c_ptr.type.element_ty) + + # Add bias, it's already in output format, so add it after conversion. + if bias_ptr: + offsets_bias = offsets_bn + bias_ptrs = bias_ptr + offsets_bias + bias_mask = offsets_bias < N + bias = tl.load(bias_ptrs, bias_mask) + c += bias + + # Save output + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64) + offs_cm = offs_cm.to(tl.int64) + offs_cn = offs_cn.to(tl.int64) + c_ptrs = c_ptr + 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) + + +# input - [M, K] +# weight - [K, N] +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py +def triton_scaled_mm( + input: torch.Tensor, + weight: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + out_dtype: type[torch.dtype], + bias: Optional[torch.Tensor] = None, + block_size_m: int = 32, + block_size_n: int = 32, + block_size_k: int = 32, + use_heuristic=True, +) -> torch.Tensor: + M, K = input.shape + N = weight.shape[1] + + assert N > 0 and K > 0 and M > 0 + assert weight.shape[0] == K + assert input.dtype == weight.dtype + + scale_a = scale_a.reshape(-1, 1) if scale_a.dim() <= 1 else scale_a + scale_b = scale_b.reshape(-1, 1) if scale_b.dim() <= 1 else scale_b + + assert scale_a.dtype == scale_b.dtype and scale_a.is_floating_point() + assert scale_a.shape[1] == 1 and (scale_a.shape[0] == 1 or scale_a.shape[0] == M) + assert scale_b.shape[1] == 1 and (scale_b.shape[0] == 1 or scale_b.shape[0] == N) + assert out_dtype.is_floating_point + assert bias is None or bias.is_floating_point() + assert is_weak_contiguous(input) + assert is_weak_contiguous(weight) + + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + + result = torch.empty((M, N), dtype=out_dtype, device=input.device) + + has_scalar = lambda x: x.shape[0] == 1 and x.shape[1] == 1 + + if use_heuristic: + is_small_N = N < 8192 + next_power_of_2_M = max(32, triton.next_power_of_2(M)) + if next_power_of_2_M <= 32: + tile_shape = (64, 64, 256) if is_small_N else (64, 128, 256) + elif next_power_of_2_M <= 64: + tile_shape = (64, 64, 256) + elif next_power_of_2_M <= 128: + tile_shape = (64, 128, 128) + else: + tile_shape = (128, 128, 128) + + block_size_m, block_size_n, block_size_k = tile_shape + + block_size_sa = 1 if has_scalar(scale_a) else block_size_m + block_size_sb = 1 if has_scalar(scale_b) else block_size_n + + accumulator_dtype = tl.float32 if input.is_floating_point() else tl.int32 + + # A = input, B = weight, C = result + # A = M x K, B = K x N, C = M x N + scaled_mm_kernel[grid]( + input, + weight, + scale_a, + scale_b, + result, + bias, + M, + N, + K, + input.stride(0), + input.stride(1), + weight.stride(0), + weight.stride(1), + result.stride(0), + result.stride(1), + accumulator_dtype, + BLOCK_SIZE_M=block_size_m, + BLOCK_SIZE_N=block_size_n, + BLOCK_SIZE_K=block_size_k, + BLOCK_SIZE_SCALE_A=block_size_sa, + BLOCK_SIZE_SCALE_B=block_size_sb, + ) + + return result.to(out_dtype) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py new file mode 100644 index 000000000..998423b86 --- /dev/null +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -0,0 +1,829 @@ +from typing import Callable, List, Optional, Tuple + +import torch + +from sglang.srt.layers.quantization import deep_gemm_wrapper +from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 +from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil +from sglang.srt.utils import is_sm100_supported + +try: + from vllm import _custom_ops as ops + + VLLM_AVAILABLE = True +except ImportError: + VLLM_AVAILABLE = False + +from sglang.srt.layers.quantization.fp8_kernel import ( + fp8_dtype, + fp8_max, + is_fp8_fnuz, + per_token_group_quant_fp8, + scaled_fp8_quant, + sglang_per_token_quant_fp8, + static_quant_fp8, + triton_scaled_mm, + w8a8_block_fp8_matmul_deepgemm, + w8a8_block_fp8_matmul_triton, +) +from sglang.srt.utils import ( + align, + ceil_div, + get_bool_env_var, + get_cuda_version, + get_device_capability, + is_cuda, + is_flashinfer_available, + is_hip, +) + +_is_hip = is_hip() +_is_cuda = is_cuda() +_is_fp8_fnuz = is_fp8_fnuz() + +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _use_aiter: + import aiter + from aiter import gemm_a8w8_blockscale, gemm_a8w8_bpreshuffle, get_hip_quant + + aiter_per1x128_quant = get_hip_quant(aiter.QuantType.per_1x128) + +if _is_cuda: + from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm + +use_vllm_cutlass_w8a8_fp8_kernel = get_bool_env_var("USE_VLLM_CUTLASS_W8A8_FP8_KERNEL") +use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL") + +# Input scaling factors are no longer optional in _scaled_mm starting +# from pytorch 2.5. Allocating a dummy tensor to pass as input_scale +TORCH_DEVICE_IDENTITY = None + + +def use_rowwise_torch_scaled_mm(): + _TORCH_VERSION = torch.__version__.split("+")[0] + try: + _TORCH_VERSION_TUPLE = tuple(map(int, _TORCH_VERSION.split(".")[:3])) + except ValueError: + _TORCH_VERSION_TUPLE = (0, 0, 0) + if _is_hip: + # The condition to determine if it is on a platform that supports + # torch._scaled_mm rowwise feature. + # The condition is determined once as the operations + # are time consuming. + return get_device_capability() >= (9, 4) and _TORCH_VERSION_TUPLE >= (2, 7, 0) + return False + + +USE_ROWWISE_TORCH_SCALED_MM = use_rowwise_torch_scaled_mm() + + +def cutlass_fp8_supported(): + if not _is_cuda: + return False + major, minor = get_device_capability() + cuda_version = get_cuda_version() + if major >= 9: + return cuda_version >= (12, 0) + elif major == 8 and minor == 9: + return cuda_version >= (12, 4) + return False + + +def normalize_e4m3fn_to_e4m3fnuz( + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + assert weight.dtype == torch.float8_e4m3fn + # The bits pattern 10000000(-128) represents zero in e4m3fn + # but NaN in e4m3fnuz. So here we set it to 0. + # https://onnx.ai/onnx/technical/float8.html + weight_as_int8 = weight.view(torch.int8) + ROCM_FP8_NAN_AS_INT = -128 + weight_as_int8[weight_as_int8 == ROCM_FP8_NAN_AS_INT] = 0 + weight = weight_as_int8.view(torch.float8_e4m3fnuz) + + # For the same bits representation, e4m3fnuz value is half of + # the e4m3fn value, so we should double the scaling factor to + # get the same dequantized value. + # https://onnx.ai/onnx/technical/float8.html + weight_scale = weight_scale * 2.0 + if input_scale is not None: + input_scale = input_scale * 2.0 + return weight, weight_scale, input_scale + + +# TODO(ch-wan): define these backends in --moe-runner-backend +def cutlass_block_fp8_supported() -> bool: + if not get_bool_env_var("SGLANG_SUPPORT_CUTLASS_BLOCK_FP8"): + return False + if _is_cuda: + major, minor = torch.cuda.get_device_capability() + sm_version = major * 10 + minor + cuda_version = tuple(map(int, torch.version.cuda.split("."))) + if cuda_version >= (12, 0) and sm_version >= 90: + return True + return False + + +CUTLASS_BLOCK_FP8_SUPPORTED = cutlass_block_fp8_supported() +ENABLE_FLASHINFER_GEMM = ( + get_bool_env_var("SGLANG_ENABLE_FLASHINFER_GEMM") + and is_sm100_supported() + and is_flashinfer_available() +) +if ENABLE_FLASHINFER_GEMM: + from flashinfer.gemm import gemm_fp8_nt_groupwise + + +def dispatch_w8a8_block_fp8_linear() -> Callable: + if ENABLE_FLASHINFER_GEMM: + return flashinfer_gemm_w8a8_block_fp8_linear + elif CUTLASS_BLOCK_FP8_SUPPORTED: + return cutlass_w8a8_block_fp8_linear_with_fallback + elif _use_aiter: + return aiter_w8a8_block_fp8_linear + elif deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + return deepgemm_w8a8_block_fp8_linear_with_fallback + else: + return triton_w8a8_block_fp8_linear + + +def flashinfer_gemm_w8a8_block_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = sglang_per_token_group_quant_fp8( + input_2d, block_size[1], column_major_scales=True + ) + # TRTLLM requires column-major scaling factors + output = gemm_fp8_nt_groupwise( + q_input, + weight, + x_scale, + weight_scale, + out_dtype=input_2d.dtype, + backend="trtllm", + ) + + if bias is not None: + output += bias + + return output.to(dtype=input_2d.dtype).view(*output_shape) + + +def cutlass_w8a8_block_fp8_linear_with_fallback( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + + # TODO: add more robust shape check here + shape_supported = weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0 + + if not shape_supported: + # fallback to triton + return triton_w8a8_block_fp8_linear( + input, weight, block_size, weight_scale, input_scale, bias + ) + + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = per_token_group_quant_fp8( + input_2d, block_size[1], column_major_scales=True + ) + output = fp8_blockwise_scaled_mm( + q_input, weight.T, x_scale, weight_scale.T, out_dtype=input_2d.dtype + ) + if bias is not None: + output += bias + return output.to(dtype=input_2d.dtype).view(*output_shape) + + +def deepgemm_w8a8_block_fp8_linear_with_fallback( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + + output_dtype = input.dtype + dtype_supported = output_dtype == torch.bfloat16 + + # TODO: https://github.com/sgl-project/sglang/pull/6890#issuecomment-2943395737 + shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0 + + if not (shape_supported and dtype_supported): + # fall back to triton + return triton_w8a8_block_fp8_linear( + input, weight, block_size, weight_scale, input_scale, bias + ) + + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = sglang_per_token_group_quant_fp8( + input_2d, + block_size[1], + column_major_scales=True, + scale_tma_aligned=True, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + + output = w8a8_block_fp8_matmul_deepgemm( + q_input, weight, x_scale, weight_scale, block_size, output_dtype=output_dtype + ) + if bias is not None: + output += bias + return output.to(dtype=output_dtype).view(*output_shape) + + +def aiter_w8a8_block_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = aiter_per1x128_quant(input_2d, quant_dtype=aiter.dtypes.fp8) + output = gemm_a8w8_blockscale( + q_input, weight, x_scale, weight_scale, dtype=input.dtype + ) + + if bias is not None: + output += bias + + return output.to(dtype=input_2d.dtype).view(*output_shape) + + +def triton_w8a8_block_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = per_token_group_quant_fp8( + input_2d, block_size[1], column_major_scales=False + ) + output = w8a8_block_fp8_matmul_triton( + q_input, weight, x_scale, weight_scale, block_size, output_dtype=input_2d.dtype + ) + if bias is not None: + output += bias + return output.to(dtype=input_2d.dtype).view(*output_shape) + + +def dequant_mxfp4( + w_block: torch.Tensor, + w_scale: torch.Tensor, + out_dtype, +) -> torch.Tensor: + """ + :param w_block: (batch, n, k, 16), uint8, pack two mxfp4 into one byte + :param w_scale: (batch, n, k), uint8 + :return: (batch, n, k * 32), float32 + """ + + assert w_block.dtype == torch.uint8 + assert w_scale.dtype == torch.uint8 + + batch, n, k, pack_dim = w_block.shape + batch_, n_, k_ = w_scale.shape + assert pack_dim == 16 + assert batch == batch_ + assert n == n_ + assert k == k_ + + out_raw = MXFP4QuantizeUtil.dequantize( + quantized_data=w_block, scale=w_scale, dtype=out_dtype, block_sizes=[32] + ) + return out_raw.reshape(batch, n, k * 32) + + +def input_to_float8( + x: torch.Tensor, dtype: torch.dtype = fp8_dtype +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function quantizes input values to float8 values with tensor-wise quantization.""" + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).float().clamp(min=1e-12) + + if _is_fp8_fnuz: + dtype = fp8_dtype + fp_max = fp8_max + else: + finfo = torch.finfo(dtype) + fp_max = finfo.max + + scale = fp_max / amax + x_scl_sat = (x.float() * scale).clamp(min=-fp_max, max=fp_max) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + + +def block_quant_to_tensor_quant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: List[int], +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function converts block-wise quantization to tensor-wise quantization. + The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale + and the block size. + The outputs are tensor-wise quantization tensor and tensor-wise quantization scale. + Note only float8 is supported for now. + """ + block_n, block_k = block_size[0], block_size[1] + n, k = x_q_block.shape + n_tiles = (n + block_n - 1) // block_n + k_tiles = (k + block_k - 1) // block_k + assert n_tiles == x_s.shape[0] + assert k_tiles == x_s.shape[1] + + x_dq_block = x_q_block.to(torch.float32) + + x_dq_block_tiles = [ + [ + x_dq_block[ + j * block_n : min((j + 1) * block_n, n), + i * block_k : min((i + 1) * block_k, k), + ] + for i in range(k_tiles) + ] + for j in range(n_tiles) + ] + + for i in range(k_tiles): + for j in range(n_tiles): + x_dq_block_tiles[j][i][:, :] = x_dq_block_tiles[j][i] * x_s[j][i] + + x_q_tensor, scale = ( + scaled_fp8_quant(x_dq_block) + if _is_cuda + else input_to_float8(x_dq_block, dtype=x_q_block.dtype) + ) + return x_q_tensor, scale + + +def block_quant_dequant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: List[int], + dtype: torch.dtype, +) -> torch.Tensor: + """This function converts block-wise quantization to unquantized. + The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale + and the block size. + The output is an unquantized tensor with dtype. + """ + block_n, block_k = block_size[0], block_size[1] + *_, n, k = x_q_block.shape + + # ... n_scale k_scale -> ... (n_scale block_n) (k_scale block_k) + x_scale_repeat = x_s.repeat_interleave(block_n, dim=-2).repeat_interleave( + block_k, dim=-1 + ) + x_scale_repeat = x_scale_repeat[..., :n, :k] + + return (x_q_block.to(torch.float32) * x_scale_repeat).to(dtype) + + +def requant_weight_ue8m0_inplace(weight, weight_scale_inv, weight_block_size): + assert isinstance(weight, torch.nn.Parameter) + assert isinstance(weight_scale_inv, torch.nn.Parameter) + weight.data, weight_scale_inv.data = _requant_weight_ue8m0( + weight, weight_scale_inv, weight_block_size + ) + + +def _requant_weight_ue8m0( + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + weight_block_size: List[int], +): + assert weight_block_size == [128, 128] + + *_, n, k = weight.shape + + weight_dequant = block_quant_dequant( + weight, + weight_scale_inv, + weight_block_size, + torch.bfloat16, + ) + + weight_dequant_flat = weight_dequant.view((-1, k)) + out_w_flat, out_s_flat = per_block_cast_to_fp8(weight_dequant_flat) + + out_w = out_w_flat.view(weight.shape) + out_s = out_s_flat.view(weight_scale_inv.shape) + + # NOTE copy and modified from DeepGEMM + def _transform_scale(sf, mn: int): + import deep_gemm.utils.layout + + sf = sf.index_select(-2, torch.arange(mn, device=sf.device) // 128) + sf = deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor(sf) + return sf + + out_s = _transform_scale(out_s, mn=out_w.shape[-2]) + + return out_w, out_s + + +# COPIED FROM DeepGEMM +def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + m, n = x.shape + x_padded = torch.zeros( + (align(m, 128), align(n, 128)), dtype=x.dtype, device=x.device + ) + x_padded[:m, :n] = x + x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) + sf = ceil_to_ue8m0(x_amax / 448.0) + x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn) + return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view( + x_view.size(0), x_view.size(2) + ) + + +# COPIED FROM DeepGEMM +def ceil_to_ue8m0(x: torch.Tensor): + return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) + + +def channel_quant_to_tensor_quant( + x_q_channel: torch.Tensor, + x_s: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + x_dq_channel = x_q_channel.to(torch.float32) * x_s + x_q_tensor, scale = ( + scaled_fp8_quant(x_dq_channel) + if _is_cuda + else input_to_float8(x_dq_channel, dtype=x_q_channel.dtype) + ) + return x_q_tensor, scale + + +def _process_scaled_mm_output(output, input_2d_shape, output_shape): + if type(output) is tuple and len(output) == 2: + output = output[0] + return torch.narrow(output, 0, 0, input_2d_shape[0]).view(*output_shape) + + +def _apply_fallback_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + input_2d_shape, + output_shape, + bias, + input_dtype, +): + global TORCH_DEVICE_IDENTITY + if TORCH_DEVICE_IDENTITY is None: + TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32, device=weight.device) + + output = torch._scaled_mm( + qinput, + weight, + scale_a=TORCH_DEVICE_IDENTITY, + scale_b=TORCH_DEVICE_IDENTITY, + out_dtype=torch.float32, + ) + + output = _process_scaled_mm_output(output, input_2d_shape, output_shape) + x_scale = torch.narrow(x_scale, 0, 0, input_2d_shape[0]) + + output = output * x_scale * weight_scale.t() + if bias is not None: + output = output + bias + return output.to(dtype=input_dtype) + + +def apply_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + input_scale_ub: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + cutlass_fp8_supported: bool = cutlass_fp8_supported(), + use_per_token_if_dynamic: bool = False, + pad_output: Optional[bool] = None, + compressed_tensor_quant: bool = False, +) -> torch.Tensor: + # Note: we pad the input because torch._scaled_mm is more performant + # for matrices with batch dimension > 16. + # This could change in the future. + # We also don't pad when using torch.compile, + # as it breaks with dynamic shapes. + if pad_output is None: + pad_output = ( + not get_bool_env_var("SGLANG_ENABLE_TORCH_COMPILE") + and not cutlass_fp8_supported + ) + output_padding = 17 if pad_output else None + + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[1]] + + if compressed_tensor_quant: + # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A + # for sgl-kernel fp8_scaled_mm, it support per channel W now + if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]: + qinput, x_scale = scaled_fp8_quant( + input_2d, + input_scale, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + + # Fused GEMM_DQ + if VLLM_AVAILABLE and use_vllm_cutlass_w8a8_fp8_kernel: + # Fall back to vllm cutlass w8a8 fp8 kernel + output = ops.cutlass_scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + ) + else: + assert ( + weight_scale.numel() == weight.shape[1] + ), "cutlass w8a8 fp8 sgl-kernel only supports per-channel scale" + + cutlass_compatible_b = ( + weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0 + ) + if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel: + # Massage the input to be 2D + qinput = qinput.view(-1, qinput.shape[-1]) + output = triton_scaled_mm( + qinput, weight, x_scale, weight_scale, input.dtype, bias + ) + else: + output = fp8_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + out_dtype=input.dtype, + bias=bias, + ) + return output.view(*output_shape) + + # torch.scaled_mm supports per tensor weights + activations only + # so fallback to naive if per channel or per token + else: + # Maybe apply padding to output, see comment in __init__ + qinput, x_scale = ( + scaled_fp8_quant( + input_2d, + input_scale, + num_token_padding=output_padding, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + if _is_cuda + else ops.scaled_fp8_quant( + input_2d, + input_scale, + num_token_padding=output_padding, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + ) + + per_tensor_weights = weight_scale.numel() == 1 + per_tensor_activations = x_scale.numel() == 1 + + if per_tensor_weights and per_tensor_activations: + # Fused GEMM_DQ + output = torch._scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + ) + return _process_scaled_mm_output(output, input_2d.shape, output_shape) + + elif ( + use_per_token_if_dynamic + and not per_tensor_weights + and not per_tensor_activations + and (USE_ROWWISE_TORCH_SCALED_MM or _use_aiter) + ): + # into this sector means use dynamic per-token-per-channel quant + # per-token scale quant for input matrix, every row(one token) have one scale factor + # per-channel scale quant for weight matrix, every col(one channel) have one scale factor + if _use_aiter: + # gemm_a8w8_bpreshuffle(XQ, WQ, x_scale, w_scale, dtype) + # XQ -> input tensor, shape = (m, k) + # WQ -> weight tensor, shape = (n, k), with preshuffe get better perf + # x_scale -> input scale tensor, shape = (m, 1) + # w_scale -> weight scale tensor, shape = (n ,1) + # dtype -> output dtype + output = gemm_a8w8_bpreshuffle( + XQ=qinput, + WQ=weight, + x_scale=x_scale, + w_scale=weight_scale, + dtype=input.dtype, + ) + if bias is not None: + output += bias + return _process_scaled_mm_output( + output, input_2d.shape, [*input.shape[:-1], weight.shape[0]] + ) + else: + # For now validated on ROCm platform + # fp8 rowwise scaling in torch._scaled_mm is introduced in + # https://github.com/pytorch/pytorch/pull/144432 using hipBLASLt + # and ROCm 6.3, which only exists in torch 2.7 and above. + # For CUDA platform please validate if the + # torch._scaled_mm support rowwise scaled GEMM + # Fused GEMM_DQ Rowwise GEMM + output = torch._scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale.t(), + bias=bias, + ) + return _process_scaled_mm_output( + output, input_2d.shape, output_shape + ) + else: + # Fallback for channelwise case, where we use unfused DQ + # due to limitations with scaled_mm + + # Symmetric quantized GEMM by definition computes the following: + # C = (s_x * X) (s_w * W) + bias + # This is equivalent to dequantizing the weights and activations + # before applying a GEMM. + # + # In order to compute quantized operands, a quantized kernel + # will rewrite the above like so: + # C = s_w * s_x * (X * W) + bias + # + # For the scaled_mm fallback case, we break this down, since it + # does not support s_w being a vector. + return _apply_fallback_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + input_2d.shape, + output_shape, + bias, + input.dtype, + ) + else: + # cutlass w8a8 fp8 sgl-kernel only supports per-token scale + if input_scale is not None: + assert input_scale.numel() == 1 + # broadcast per-tensor scale to per-token scale when supporting cutlass + qinput, x_scale = static_quant_fp8( + input_2d, input_scale, repeat_scale=cutlass_fp8_supported + ) + else: + # default use per-token quantization if dynamic + if _is_cuda: + qinput, x_scale = sglang_per_token_quant_fp8(input_2d) + else: + # TODO(kkhuang): temporarily enforce per-tensor activation scaling if weight is per-tensor scaling + # final solution should be: 1. add support to per-tensor activation scaling. + # 2. solve the torch.compile error from weight_scale.numel() == 1 and x_scale.numel() > 1 (below line#308) + if _is_hip and weight_scale.numel() == 1: + qinput, x_scale = ops.scaled_fp8_quant( + input_2d, + input_scale, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + else: + qinput, x_scale = per_token_group_quant_fp8( + input_2d, group_size=input_2d.shape[1] + ) + + if cutlass_fp8_supported: + try: + if VLLM_AVAILABLE and use_vllm_cutlass_w8a8_fp8_kernel: + # Fall back to vllm cutlass w8a8 fp8 kernel + output = ops.cutlass_scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + ) + else: + assert ( + weight_scale.numel() == weight.shape[1] + ), "cutlass w8a8 fp8 sgl-kernel only supports per-channel scale" + + cutlass_compatible_b = ( + weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0 + ) + if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel: + # Massage the input to be 2D + qinput = qinput.view(-1, qinput.shape[-1]) + output = triton_scaled_mm( + qinput, weight, x_scale, weight_scale, input.dtype, bias + ) + else: + output = fp8_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + out_dtype=input.dtype, + bias=bias, + ) + return output.view(*output_shape) + except (ImportError, NameError, AttributeError): + pass + + # torch.scaled_mm supports per tensor weights + activations only + # so fallback to naive if per channel or per token + per_tensor_weights = weight_scale.numel() == 1 + per_tensor_activations = x_scale.numel() == 1 + + if per_tensor_weights and per_tensor_activations: + # Fused GEMM_DQ + output = torch._scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + ) + return _process_scaled_mm_output(output, input_2d.shape, output_shape) + + else: + # Fallback for channelwise case, where we use unfused DQ + # due to limitations with scaled_mm + + # Symmetric quantized GEMM by definition computes the following: + # C = (s_x * X) (s_w * W) + bias + # This is equivalent to dequantizing the weights and activations + # before applying a GEMM. + # + # In order to compute quantized operands, a quantized kernel + # will rewrite the above like so: + # C = s_w * s_x * (X * W) + bias + # + # For the scaled_mm fallback case, we break this down, since it + # does not support s_w being a vector. + return _apply_fallback_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + input_2d.shape, + output_shape, + bias, + input.dtype, + ) + + +def can_auto_enable_marlin_fp8() -> bool: + try: + major, minor = get_device_capability() + sm = major * 10 + minor + return 80 <= sm < 89 + except Exception: + return False diff --git a/python/sglang/srt/layers/quantization/fpgemm_fp8.py b/python/sglang/srt/layers/quantization/fpgemm_fp8.py new file mode 100644 index 000000000..5a78626ff --- /dev/null +++ b/python/sglang/srt/layers/quantization/fpgemm_fp8.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import logging +from typing import Any, Optional + +import torch +from torch.nn import Module +from torch.nn.parameter import Parameter + +from sglang.srt.layers.linear import LinearBase +from sglang.srt.layers.parameter import ChannelQuantScaleParameter, ModelWeightParameter +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + can_auto_enable_marlin_fp8, + cutlass_fp8_supported, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.layers.quantization.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + prepare_fp8_layer_for_marlin, +) +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.layers.quantization.utils import is_layer_skipped, replace_parameter +from sglang.srt.utils import get_bool_env_var, is_cuda + +_is_cuda = is_cuda() +_is_fp8_fnuz = is_fp8_fnuz() + +logger = logging.getLogger(__name__) + + +class FBGEMMFp8Config(QuantizationConfig): + """Config class for FBGEMM Fp8.""" + + def __init__(self, ignore_list: list[str], input_scale_ub: float): + super().__init__() + self.ignore_list = ignore_list if ignore_list else [] + self.input_scale_ub = input_scale_ub + + # For GPUs that lack FP8 hardware suspport, we can leverage the Marlin + # kernel for fast weight-only FP8 quantization + # self.use_marlin = not marlin_fp8_supported() + self.use_marlin = False + if _is_cuda: + force_marlin = get_bool_env_var("SGLANG_FORCE_FP8_MARLIN") + auto_enable = can_auto_enable_marlin_fp8() + self.use_marlin = force_marlin or auto_enable + + @classmethod + def get_name(cls) -> str: + return "fbgemm_fp8" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> FBGEMMFp8Config: + ignore_list = cls.get_from_keys(config, ["modules_to_not_convert"]) + input_scale_ub = cls.get_from_keys(config, ["activation_scale_ub"]) + return cls(ignore_list=ignore_list, input_scale_ub=input_scale_ub) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignore_list, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + return FBGEMMFp8LinearMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class FBGEMMFp8LinearMethod(LinearMethodBase): + + def __init__(self, quant_config: FBGEMMFp8Config): + self.quant_config = quant_config + # self.fp8_linear = Fp8LinearOp( + # act_quant_static=False, act_quant_group_shape=GroupShape.PER_TOKEN) + self.out_dtype = torch.get_default_dtype() + self.cutlass_fp8_supported = cutlass_fp8_supported() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + # maybe_create_device_identity() + weight_loader = extra_weight_attrs.get("weight_loader") + del input_size, output_size + output_size_per_partition = sum(output_partition_sizes) + + layer.logical_widths = output_partition_sizes + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # WEIGHT + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=torch.float8_e4m3fn, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + weight_scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale", weight_scale) + + # INPUT SCALE UPPER BOUND + input_scale_ub = torch.nn.Parameter( + torch.tensor((self.quant_config.input_scale_ub), dtype=torch.float32), + requires_grad=False, + ) + layer.input_scale_ub = input_scale_ub + + def process_weights_after_loading(self, layer: Module) -> None: + # required by torch.compile + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) + layer.weight = Parameter(layer.weight.data, requires_grad=False) + + weight = layer.weight + + if _is_fp8_fnuz: + weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, weight_scale=layer.weight_scale, input_scale=None + ) + if input_scale is not None: + layer.input_scale = Parameter(input_scale, requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + + layer.weight = Parameter(weight.t(), requires_grad=False) + if self.quant_config.use_marlin: + prepare_fp8_layer_for_marlin(layer) + # Activations not quantized for marlin. + del layer.input_scale_ub + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + if self.quant_config.use_marlin: + return apply_fp8_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) + + return apply_fp8_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=None, + input_scale_ub=layer.input_scale_ub, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + use_per_token_if_dynamic=False, + ) diff --git a/python/sglang/srt/layers/quantization/gptq.py b/python/sglang/srt/layers/quantization/gptq.py new file mode 100644 index 000000000..ccd3d46f7 --- /dev/null +++ b/python/sglang/srt/layers/quantization/gptq.py @@ -0,0 +1,1097 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from fractions import Fraction +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union + +import torch + +from sglang.srt.layers.parameter import ( + BasevLLMParameter, + ChannelQuantScaleParameter, + GroupQuantScaleParameter, + PackedColumnParameter, + PackedvLLMParameter, + RowvLLMParameter, + permute_param_layout_, +) +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.marlin_utils import ( + apply_gptq_marlin_linear, + check_marlin_supported, + check_marlin_supports_shape, + marlin_is_k_full, + marlin_make_empty_g_idx, + marlin_make_workspace, + marlin_moe_permute_scales, + marlin_permute_scales, + marlin_repeat_scales_on_all_ranks, + marlin_sort_g_idx, + marlin_zero_points, + verify_marlin_supported, +) +from sglang.srt.layers.quantization.utils import ( + get_linear_quant_method, + get_scalar_types, + replace_parameter, + unpack_cols, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig + from sglang.srt.layers.moe.token_dispatcher import ( + StandardDispatchOutput, + CombineInput, + ) + +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() + +if _is_cuda: + from sgl_kernel import fused_marlin_moe, gptq_gemm, gptq_marlin_repack, gptq_shuffle + + +logger = logging.getLogger(__name__) +ScalarType, scalar_types = get_scalar_types() + + +def check_marlin_format(hf_quant_cfg: Dict[str, Any]) -> bool: + # compat: gptqmodel and autogptq (eol) main use checkpoint_format: str + # compat: autogptq <=0.7.1 is_marlin_format: bool + return hf_quant_cfg.get("checkpoint_format") == "marlin" or hf_quant_cfg.get( + "is_marlin_format", False + ) + + +def gptq_marlin_moe_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, +) -> torch.Tensor: + num_experts = b_q_weight.shape[0] + assert size_k % 16 == 0 + output = torch.empty( + (num_experts, size_k // 16, size_n * (num_bits // 2)), + device=b_q_weight.device, + dtype=b_q_weight.dtype, + ) + for e in range(num_experts): + output[e] = gptq_marlin_repack(b_q_weight[e], perm[e], size_k, size_n, num_bits) + return output + + +@dataclass +class MarlinLinearLayerConfig: + full_weight_shape: tuple[int, int] # [in, out] + partition_weight_shape: tuple[int, int] + weight_type: ScalarType + act_type: torch.dtype + group_size: int + zero_points: bool + has_g_idx: bool + + +class GPTQConfig(QuantizationConfig): + """Config class for GPTQ. + + Reference: https://arxiv.org/abs/2210.17323 + """ + + def __init__( + self, + weight_bits: int, + group_size: int, + desc_act: bool, + lm_head_quantized: bool, + dynamic: Dict[str, Dict[str, Union[int, bool]]], + ) -> None: + # GPTQModel use `dynamic` config property to allow per module + # quantization config so each module can be individually optimized. + # Format is Dict[str, Dict] where key is a regex string that can + # perform both positive ("+:" prefixed) or negative ("-:" prefixed) + # matching of a module. + # Default to positive match, override base quant config mode, if no + # prefix is used. Value is in dict format of field key and override + # value. + # Negative matching will skip quantization init for this module + # entirely: + # non-quantized inference. More details and quantization examples can be + # found at: https://github.com/ModelCloud/GPTQModel + # Example: + # # last 1/2 of the layers 10-21 has 8bit vs 4bit for 0-9 + # # last 1/4 of the layers 16-21 has 8bit and group_size 64 + # dynamic = { + # #`.*\.` matches the layers_node prefix + # # positive match layer 10-15 + # r"+:.*\.(?:1[0-5])\..*": {"bits": 8,}, + # # positive match layer 16-21 + # r"+:.*\.(?:1[6-9]|20|21)\..*": {"bits": 8, "group_size": 64,}, + # r"-:.*\.moe\..*": {}, # negative match (skip) all `moe` layers + # } + super().__init__() + self.dynamic = dynamic + + self.weight_bits = weight_bits + self.group_size = group_size + self.desc_act = desc_act + self.lm_head_quantized = lm_head_quantized + self.pack_factor = Fraction(32, self.weight_bits) + if self.weight_bits not in [2, 3, 4, 8]: + raise ValueError( + "Currently, only 2/3/4/8-bit weight quantization is " + f"supported for GPTQ, but got {self.weight_bits} bits." + ) + + def __repr__(self) -> str: + return ( + f"GPTQConfig(weight_bits={self.weight_bits}, " + f"group_size={self.group_size}, " + f"desc_act={self.desc_act})," + f"lm_head_quantized={self.lm_head_quantized}), " + f"dynamic={self.dynamic}" + ) + + def get_scaled_act_names(self) -> List[str]: + """Returns the activation function names that should be post-scaled. + + For now, this is only used by AWQ. + """ + raise NotImplementedError + + @classmethod + def get_name(cls) -> str: + return "gptq" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.half] + + @classmethod + # Need to figure it out + def get_min_capability(cls) -> int: + return 60 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["quantize_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> GPTQConfig: + dynamic = cls.get_from_keys_or(config, ["dynamic"], default={}) + dynamic = {} if dynamic is None else dynamic + + weight_bits = cls.get_from_keys(config, ["bits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + desc_act = cls.get_from_keys(config, ["desc_act"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + return cls(weight_bits, group_size, desc_act, lm_head_quantized, dynamic) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[LinearMethodBase]: + # Delay the import to avoid circular dependency + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, FusedMoE): + raise TypeError("GPTQ Method does not support MoE, please use gptq_marlin") + else: + return get_linear_quant_method( + self, layer, prefix=prefix, linear_method_cls=GPTQLinearMethod + ) + + +class GPTQMarlinConfig(QuantizationConfig): + """Config class for GPTQ Marlin""" + + # (num_bits, is_sym) -> quant_type + TYPE_MAP = { + (4, True): scalar_types.uint4b8, + (8, True): scalar_types.uint8b128, + } + + def __init__( + self, + weight_bits: int, + group_size: int, + desc_act: bool, + is_sym: bool, + lm_head_quantized: bool, + dynamic: Dict[str, Dict[str, Union[int, bool]]], + full_config: Dict[str, Any], + ) -> None: + super().__init__() + if desc_act and group_size == -1: + # In this case, act_order == True is the same as act_order == False + # (since we have only one group per output channel) + desc_act = False + + # GPTQModel use `dynamic` config property to allow per module + # quantization config so each module can be individually optimized. + # Format is Dict[str, Dict] where key is a regex string that can + # perform both positive ("+:" prefixed) or negative ("-:" prefixed) + # matching of a module. + # Default to positive match, override base quant config mode, if no + # prefix is used. Value is in dict format of field key and override + # value. + # Negative matching will skip quantization init for this module + # entirely: + # non-quantized inference. More details and quantization examples can be + # found at: https://github.com/ModelCloud/GPTQModel + # Example: + # # last 1/2 of the layers 10-21 has 8bit vs 4bit for 0-9 + # # last 1/4 of the layers 16-21 has 8bit and group_size 64 + # dynamic = { + # #`.*\.` matches the layers_node prefix + # # positive match layer 10-15 + # r"+:.*\.(?:1[0-5])\..*": {"bits": 8,}, + # # positive match layer 16-21 + # r"+:.*\.(?:1[6-9]|20|21)\..*": {"bits": 8, "group_size": 64,}, + # r"-:.*\.moe\..*": {}, # negative match (skip) all `moe` layers + # } + self.dynamic = dynamic + + self.weight_bits = weight_bits + self.is_sym = is_sym + + self.pack_factor = 32 // weight_bits # packed into int32 + self.group_size = group_size + self.desc_act = desc_act + self.lm_head_quantized = lm_head_quantized + self.full_config = full_config + + if (weight_bits, is_sym) not in self.TYPE_MAP: + raise ValueError( + "Unsupported quantization config: " f"bits={weight_bits}, sym={is_sym}" + ) + + # (num_bits, is_sym) -> quant_type + self.quant_type = self.TYPE_MAP[(weight_bits, is_sym)] + + def __repr__(self) -> str: + return ( + f"GPTQMarlinConfig(quant_type={self.quant_type}, " + f"group_size={self.group_size}, " + f"desc_act={self.desc_act}, " + f"lm_head_quantized={self.lm_head_quantized}), " + f"dynamic={self.dynamic}" + ) + + def get_scaled_act_names(self) -> List[str]: + """Returns the activation function names that should be post-scaled. + + For now, this is only used by AWQ. + """ + raise NotImplementedError + + @classmethod + def get_name(cls) -> str: + return "gptq_marlin" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.half, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["quantize_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> GPTQMarlinConfig: + dynamic = cls.get_from_keys_or(config, ["dynamic"], default={}) + dynamic = {} if dynamic is None else dynamic + + weight_bits = cls.get_from_keys(config, ["bits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + desc_act = cls.get_from_keys(config, ["desc_act"]) + is_sym = cls.get_from_keys(config, ["sym"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + return cls( + weight_bits, + group_size, + desc_act, + is_sym, + lm_head_quantized, + dynamic, + config, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + is_marlin_format = check_marlin_format(hf_quant_cfg) + + can_convert = cls.is_gptq_marlin_compatible(hf_quant_cfg) + + is_valid_user_quant = ( + user_quant is None or user_quant == "marlin" or user_quant == "gptq_marlin" + ) + + if not is_marlin_format and can_convert and is_valid_user_quant: + msg = ( + "The model is convertible to {} during runtime." + " Using {} kernel.".format(cls.get_name(), cls.get_name()) + ) + logger.info(msg) + return cls.get_name() + + if not is_marlin_format and can_convert and user_quant == "gptq": + logger.info( + "Detected that the model can run with gptq_marlin" + ", however you specified quantization=gptq explicitly," + " so forcing gptq. Use quantization=gptq_marlin for" + " faster inference" + ) + return None + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + # Delay the import to avoid circular dependency + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, FusedMoE): + return GPTQMarlinMoEMethod(self) + return get_linear_quant_method(self, layer, prefix, GPTQMarlinLinearMethod) + + @classmethod + def is_gptq_marlin_compatible(cls, quant_config: Dict[str, Any]): + quant_method = quant_config.get("quant_method", "").lower() + num_bits = quant_config.get("bits") + group_size = quant_config.get("group_size") + sym = quant_config.get("sym") + desc_act = quant_config.get("desc_act") + + if not _is_cuda: + return False + + if quant_method != "gptq": + return False + + # Marlin conversion is only valid if required properties are found + if num_bits is None or group_size is None or sym is None or desc_act is None: + return False + + if (num_bits, sym) not in cls.TYPE_MAP: + return False + + return check_marlin_supported( + quant_type=cls.TYPE_MAP[(num_bits, sym)], group_size=group_size + ) + + +class GPTQLinearMethod(LinearMethodBase): + """Linear method for GPTQ. + + Args: + quant_config: The GPTQ quantization config. + """ + + def __init__(self, quant_config: GPTQConfig): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + del output_size # Unused. + weight_loader = extra_weight_attrs.get("weight_loader") + if input_size_per_partition % self.quant_config.group_size != 0: + raise ValueError( + "The input size is not aligned with the quantized " + "weight shape. This can be caused by too large " + "tensor parallel size." + ) + output_size_per_partition = sum(output_partition_sizes) + if output_size_per_partition % self.quant_config.pack_factor.numerator != 0: + raise ValueError( + "The output size is not aligned with the quantized " + "weight shape. This can be caused by too large " + "tensor parallel size." + ) + + if self.quant_config.group_size != -1: + group_size = self.quant_config.group_size + else: + group_size = input_size + + self.use_shuffle = True + scale_and_zero_size = input_size // group_size + scale_and_zero_input_dim = None + if ( + input_size != input_size_per_partition + and self.quant_config.group_size != -1 + ): + if self.quant_config.desc_act: + self.use_shuffle = False + else: + # we need to partition qzeros and scales for exllama kernel + scale_and_zero_size = input_size_per_partition // group_size + scale_and_zero_input_dim = 0 + + qweight = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition // self.quant_config.pack_factor, + output_size_per_partition, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=0, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + g_idx = RowvLLMParameter( + data=torch.tensor( + [ + i // self.quant_config.group_size + for i in range(input_size_per_partition) + ], + dtype=torch.int32, + ), + input_dim=0, + weight_loader=weight_loader, + ) + qzeros_args = { + "data": torch.empty( + scale_and_zero_size, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + "weight_loader": weight_loader, + } + weight_scale_args = { + "data": torch.empty( + scale_and_zero_size, + output_size_per_partition, + dtype=params_dtype, + ), + "weight_loader": weight_loader, + } + if scale_and_zero_input_dim is None: + scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args) + qzeros = PackedColumnParameter( + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + **qzeros_args, + ) + + else: + scales = GroupQuantScaleParameter( + output_dim=1, input_dim=0, **weight_scale_args + ) + qzeros = PackedvLLMParameter( + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + **qzeros_args, + ) + + layer.register_parameter("qweight", qweight) + layer.register_parameter("g_idx", g_idx) + layer.register_parameter("qzeros", qzeros) + layer.register_parameter("scales", scales) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # for torch.compile + layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) + layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) + layer.g_idx = torch.nn.Parameter(layer.g_idx.data, requires_grad=False) + layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) + + # exllama needs to shuffle the weight after the weight is loaded + # here we do the shuffle on first forward pass + if self.use_shuffle: + if self.quant_config.desc_act: + layer.g_idx.data = torch.argsort(layer.g_idx).to(torch.int) + else: + layer.g_idx.data = torch.empty( + (0,), dtype=torch.int, device=layer.g_idx.device + ) + gptq_shuffle(layer.qweight, layer.g_idx, self.quant_config.weight_bits) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + out_shape = x.shape[:-1] + (layer.qweight.shape[-1],) + reshaped_x = x.reshape(-1, x.shape[-1]) + + output = gptq_gemm( + reshaped_x, + layer.qweight, + layer.qzeros, + layer.scales, + layer.g_idx, + self.use_shuffle, + self.quant_config.weight_bits, + ) + if bias is not None: + output.add_(bias) + return output.reshape(out_shape) + + +class GPTQMarlinLinearMethod(LinearMethodBase): + """Linear method for GPTQ Marlin. + + Args: + quant_config: The GPTQ Marlin quantization config. + """ + + _kernel_backends_being_used: set[str] = set() + + def __init__(self, quant_config: GPTQMarlinConfig) -> None: + self.quant_config = quant_config + + # Verify supported on platform. + verify_marlin_supported( + quant_type=self.quant_config.quant_type, + group_size=self.quant_config.group_size, + ) + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + is_row_parallel = input_size != input_size_per_partition + weight_loader = extra_weight_attrs.get("weight_loader") + + self.kernel_config = MarlinLinearLayerConfig( + full_weight_shape=(input_size, output_size), + partition_weight_shape=( + input_size_per_partition, + output_size_per_partition, + ), + weight_type=self.quant_config.quant_type, + act_type=params_dtype, + group_size=self.quant_config.group_size, + zero_points=False, + has_g_idx=self.quant_config.desc_act, + ) + # Normalize group_size + if self.quant_config.group_size != -1: + group_size = self.quant_config.group_size + else: + group_size = input_size + + # Determine sharding + if marlin_repeat_scales_on_all_ranks( + self.quant_config.desc_act, self.quant_config.group_size, is_row_parallel + ): + # By setting scale_dim == None, weight_loader will + # repeat the scales on each GPU in TP>1 case. + scales_and_zp_input_dim = None + scales_and_zp_size = input_size // group_size + else: + # By setting scale_dim == 0, weight_loader will + # shard the scales in TP>1 case. + scales_and_zp_input_dim = 0 + scales_and_zp_size = input_size_per_partition // group_size + + # Quantized weights + qweight = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition // self.quant_config.pack_factor, + output_size_per_partition, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=0, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + # Activation order + g_idx = RowvLLMParameter( + data=torch.empty( + input_size_per_partition, + dtype=torch.int32, + ), + input_dim=0, + weight_loader=weight_loader, + ) + + qzeros_args = { + "data": torch.empty( + scales_and_zp_size, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + "weight_loader": weight_loader, + } + weight_scale_args = { + "data": torch.empty( + scales_and_zp_size, + output_size_per_partition, + dtype=params_dtype, + ), + "weight_loader": weight_loader, + } + + if scales_and_zp_input_dim is None: + scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args) + qzeros = PackedColumnParameter( + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + **qzeros_args, + ) + + else: + scales = GroupQuantScaleParameter( + output_dim=1, input_dim=0, **weight_scale_args + ) + qzeros = PackedvLLMParameter( + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + **qzeros_args, + ) + + layer.register_parameter("qweight", qweight) + layer.register_parameter("g_idx", g_idx) + layer.register_parameter("scales", scales) + layer.register_parameter("qzeros", qzeros) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + device = getattr(layer, "qweight").device + c = self.kernel_config + + check_marlin_supports_shape( + c.partition_weight_shape[1], # out_features + c.partition_weight_shape[0], # in_features + c.full_weight_shape[0], # in_features + c.group_size, + ) + + row_parallel = c.partition_weight_shape[0] != c.full_weight_shape[0] + self.is_k_full = marlin_is_k_full(c.has_g_idx, row_parallel) + + # Allocate marlin workspace. + self.workspace = marlin_make_workspace(device) + + # Default names since marlin requires empty parameters for these, + # TODO: remove this requirement from marlin (allow optional tensors) + self.w_q_name = "qweight" + self.w_s_name = "scales" + self.w_zp_name = "qzeros" + self.w_gidx_name = "g_idx" + + def _transform_param( + layer: torch.nn.Module, name: Optional[str], fn: Callable + ) -> None: + if name is not None and getattr(layer, name, None) is not None: + + old_param = getattr(layer, name) + new_param = fn(old_param) + # replace the parameter with torch.nn.Parameter for TorchDynamo + # compatibility + replace_parameter( + layer, name, torch.nn.Parameter(new_param.data, requires_grad=False) + ) + + def transform_w_q(x): + assert isinstance(x, BasevLLMParameter) + permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) + x.data = gptq_marlin_repack( + x.data.contiguous(), + perm=layer.g_idx_sort_indices, + size_k=c.partition_weight_shape[0], + size_n=c.partition_weight_shape[1], + num_bits=c.weight_type.size_bits, + ) + return x + + def transform_w_s(x): + assert isinstance(x, BasevLLMParameter) + permute_param_layout_(x, input_dim=0, output_dim=1) + x.data = marlin_permute_scales( + x.data.contiguous(), + size_k=c.partition_weight_shape[0], + size_n=c.partition_weight_shape[1], + group_size=c.group_size, + ) + return x + + if c.has_g_idx: + g_idx, g_idx_sort_indices = marlin_sort_g_idx( + getattr(layer, self.w_gidx_name) + ) + _transform_param(layer, self.w_gidx_name, lambda _: g_idx) + layer.g_idx_sort_indices = g_idx_sort_indices + else: + setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device)) + layer.g_idx_sort_indices = marlin_make_empty_g_idx(device) + + if c.zero_points: + grouped_k = ( + c.partition_weight_shape[0] // c.group_size if c.group_size != -1 else 1 + ) + _transform_param( + layer, + self.w_zp_name, + lambda x: marlin_zero_points( + unpack_cols( + x.t(), + c.weight_type.size_bits, + grouped_k, + c.partition_weight_shape[1], + ), + size_k=grouped_k, + size_n=c.partition_weight_shape[1], + num_bits=c.weight_type.size_bits, + ), + ) + else: + setattr(layer, self.w_zp_name, marlin_make_empty_g_idx(device)) + _transform_param(layer, self.w_q_name, transform_w_q) + _transform_param(layer, self.w_s_name, transform_w_s) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + c = self.kernel_config + + def _get_weight_params( + layer: torch.nn.Module, + ) -> tuple[ + torch.Tensor, # w_q + torch.Tensor, # w_s + Optional[torch.Tensor], # w_zp, + Optional[torch.Tensor], # w_gidx + ]: + return ( + getattr(layer, self.w_q_name), + getattr(layer, self.w_s_name), + getattr(layer, self.w_zp_name or "", None), + getattr(layer, self.w_gidx_name or "", None), + ) + + w_q, w_s, w_zp, w_gidx = _get_weight_params(layer) + + # `process_weights_after_loading` will ensure w_zp and w_gidx are not + # None for marlin + return apply_gptq_marlin_linear( + input=x, + weight=w_q, + weight_scale=w_s, + weight_zp=w_zp, # type: ignore + g_idx=w_gidx, # type: ignore + g_idx_sort_indices=layer.g_idx_sort_indices, + workspace=self.workspace, + wtype=c.weight_type, + input_size_per_partition=c.partition_weight_shape[0], + output_size_per_partition=c.partition_weight_shape[1], + is_k_full=self.is_k_full, + bias=bias, + ) + + +class GPTQMarlinMoEMethod(FusedMoEMethodBase): + """MoE Marlin method with quantization.""" + + def __init__(self, quant_config: GPTQMarlinConfig) -> None: + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + # Delay the import to avoid circular dependency + from sglang.srt.layers.linear import set_weight_attrs + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + self.is_k_full = (not self.quant_config.desc_act) or layer.moe_tp_size == 1 + + if self.quant_config.group_size != -1: + scales_size13 = hidden_size // self.quant_config.group_size + if self.quant_config.desc_act: + w2_scales_size = intermediate_size_per_partition + else: + w2_scales_size = intermediate_size_per_partition * layer.moe_tp_size + scales_size2 = w2_scales_size // self.quant_config.group_size + strategy = FusedMoeWeightScaleSupported.GROUP.value + else: + scales_size13 = 1 + scales_size2 = 1 + strategy = FusedMoeWeightScaleSupported.CHANNEL.value + + extra_weight_attrs.update({"quant_method": strategy, "is_transposed": True}) + # Fused gate_up_proj (column parallel) + w13_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size // self.quant_config.pack_factor, + 2 * intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qweight", w13_qweight) + set_weight_attrs(w13_qweight, extra_weight_attrs) + # down_proj (row parallel) + w2_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition // self.quant_config.pack_factor, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qweight", w2_qweight) + set_weight_attrs(w2_qweight, extra_weight_attrs) + # up_proj scales + w13_scales = torch.nn.Parameter( + torch.empty( + num_experts, + scales_size13, + 2 * intermediate_size_per_partition, + dtype=torch.half, + ), + requires_grad=False, + ) + layer.register_parameter("w13_scales", w13_scales) + set_weight_attrs(w13_scales, extra_weight_attrs) + # down_proj scales + w2_scales = torch.nn.Parameter( + torch.empty(num_experts, scales_size2, hidden_size, dtype=torch.half), + requires_grad=False, + ) + layer.register_parameter("w2_scales", w2_scales) + set_weight_attrs(w2_scales, extra_weight_attrs) + # dont shard the w2 scales when running act order + set_weight_attrs(w2_scales, {"load_full_w2": self.quant_config.desc_act}) + # up_proj scales + w13_qzeros = torch.nn.Parameter( + torch.empty( + num_experts, + scales_size13, + 2 * intermediate_size_per_partition // self.quant_config.pack_factor, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qzeros", w13_qzeros) + set_weight_attrs(w13_qzeros, extra_weight_attrs) + # down_proj scales + w2_qzeros = torch.nn.Parameter( + torch.empty( + num_experts, + scales_size2, + hidden_size // self.quant_config.pack_factor, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qzeros", w2_qzeros) + set_weight_attrs(w2_qzeros, extra_weight_attrs) + # dont shard the w2 scales when running act order + set_weight_attrs(w2_qzeros, {"load_full_w2": self.quant_config.desc_act}) + w13_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_g_idx", w13_g_idx) + set_weight_attrs(w13_g_idx, extra_weight_attrs) + w2_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_g_idx", w2_g_idx) + set_weight_attrs(w2_g_idx, extra_weight_attrs) + w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) + set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) + w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) + set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + + # Process act_order + if self.quant_config.desc_act: + # Get sorting based on g_idx + num_experts = layer.w13_g_idx.shape[0] + w13_g_idx_sort_indices = torch.empty_like(layer.w13_g_idx) + w2_g_idx_sort_indices = torch.empty_like(layer.w2_g_idx) + w13_sorted_g_idx = torch.empty_like(layer.w13_g_idx) + w2_sorted_g_idx = torch.empty_like(layer.w2_g_idx) + for e in range(num_experts): + w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_g_idx[e]).to( + torch.int32 + ) + w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_g_idx[e]).to( + torch.int32 + ) + w13_sorted_g_idx[e] = layer.w13_g_idx[e][w13_g_idx_sort_indices[e]] + w2_sorted_g_idx[e] = layer.w2_g_idx[e][w2_g_idx_sort_indices[e]] + replace_parameter(layer, "w13_g_idx", w13_sorted_g_idx) + replace_parameter(layer, "w2_g_idx", w2_sorted_g_idx) + replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) + else: + # Reset g_idx related tensors + num_experts = layer.w13_g_idx.shape[0] + device = layer.w13_g_idx.device + layer.w13_g_idx = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_g_idx = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + # Repack weights + marlin_w13_qweight = gptq_marlin_moe_repack( + layer.w13_qweight, + layer.w13_g_idx_sort_indices, + layer.w13_qweight.shape[1] * self.quant_config.pack_factor, + layer.w13_qweight.shape[2], + self.quant_config.weight_bits, + ) + replace_parameter(layer, "w13_qweight", marlin_w13_qweight) + marlin_w2_qweight = gptq_marlin_moe_repack( + layer.w2_qweight, + layer.w2_g_idx_sort_indices, + layer.w2_qweight.shape[1] * self.quant_config.pack_factor, + layer.w2_qweight.shape[2], + self.quant_config.weight_bits, + ) + replace_parameter(layer, "w2_qweight", marlin_w2_qweight) + # Repack scales + marlin_w13_scales = marlin_moe_permute_scales( + s=layer.w13_scales, + size_k=layer.intermediate_size_per_partition, + size_n=layer.w13_scales.shape[2], + group_size=self.quant_config.group_size, + ) + replace_parameter(layer, "w13_scales", marlin_w13_scales) + marlin_w2_scales = marlin_moe_permute_scales( + s=layer.w2_scales, + size_k=layer.w2_scales.shape[1] + * ( + self.quant_config.group_size + if self.quant_config.group_size != -1 + else self.quant_config.pack_factor + ), + size_n=layer.w2_scales.shape[2], + group_size=self.quant_config.group_size, + ) + replace_parameter(layer, "w2_scales", marlin_w2_scales) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + # Delay the import to avoid circular dependency + + assert ( + self.moe_runner_config.activation == "silu" + ), "Only SiLU activation is supported." + + # The input must currently be float16 + orig_dtype = x.dtype + x = x.half() + + topk_weights, topk_ids, router_logits = topk_output + + output = fused_marlin_moe( + x, + layer.w13_qweight, + layer.w2_qweight, + layer.w13_scales, + layer.w2_scales, + router_logits, + topk_weights, + topk_ids, + g_idx1=layer.w13_g_idx, + g_idx2=layer.w2_g_idx, + sort_indices1=layer.w13_g_idx_sort_indices, + sort_indices2=layer.w2_g_idx_sort_indices, + num_bits=self.quant_config.weight_bits, + is_k_full=self.is_k_full, + ).to(orig_dtype) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/int8_kernel.py b/python/sglang/srt/layers/quantization/int8_kernel.py new file mode 100644 index 000000000..7c6c3dbd4 --- /dev/null +++ b/python/sglang/srt/layers/quantization/int8_kernel.py @@ -0,0 +1,425 @@ +import functools +import json +import logging +import os +from typing import Any, Dict, List, Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.utils import get_device_name, is_cuda + +_is_cuda = is_cuda() +if _is_cuda: + from sgl_kernel import sgl_per_token_group_quant_int8 + +logger = logging.getLogger(__name__) + + +@triton.jit +def _per_token_quant_int8( + x_ptr, + xq_ptr, + scale_ptr, + x_sum_ptr, + stride_x, + stride_xq, + N, + CAL_SUM: tl.constexpr, + BLOCK: tl.constexpr, +): + # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282 + row_id = tl.program_id(0) + + cols = tl.arange(0, BLOCK) + mask = cols < N + + x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask, other=0.0).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10) + scale_x = absmax / 127 + x_q = x * (127 / absmax) + x_q = tl.extra.cuda.libdevice.round(x_q).to(tl.int8) + if CAL_SUM: + x_sum = tl.sum(x, axis=0) + tl.store(x_sum_ptr + row_id, x_sum.to(x_sum_ptr.dtype.element_ty)) + + tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask) + tl.store(scale_ptr + row_id, scale_x.to(scale_ptr.dtype.element_ty)) + + +def per_token_quant_int8(x, scale_dtype=torch.float32, cal_sum=False): + M = x.numel() // x.shape[-1] + N = x.shape[-1] + x_q = torch.empty_like(x, device=x.device, dtype=torch.int8) + scales = torch.empty(x.shape[:-1] + (1,), device=x.device, dtype=scale_dtype) + if cal_sum: + x_sum = torch.empty(x.shape[:-1], device=x.device, dtype=x.dtype) + else: + x_sum = None + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + + assert x.is_contiguous() + _per_token_quant_int8[(M,)]( + x, + x_q, + scales, + x_sum, + stride_x=x.stride(-2), + stride_xq=x_q.stride(-2), + N=N, + CAL_SUM=cal_sum, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=1, + ) + if cal_sum: + return x_q, scales, x_sum + else: + return x_q, scales + + +@triton.jit +def _per_token_group_quant_int8( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + # Stride of input + y_stride, + # Columns of input + N, + # Avoid to divide zero + eps, + # Information for int8 + int8_min, + int8_max, + # Meta-parameters + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group quantization on a + tensor. + + This function converts the tensor values into int8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * y_stride + y_q_ptr += g_id * y_stride + y_s_ptr += g_id + + cols = tl.arange(0, BLOCK) # N <= BLOCK + mask = cols < N + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / int8_max + y_q = tl.clamp(y / y_s, int8_min, int8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + + +def per_token_group_quant_int8( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: torch.dtype = torch.int8, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Function to perform per-token-group quantization on an input tensor `x`. + + It converts the tensor values into signed int8 values and returns the + quantized tensor along with the scaling factor used for quantization. + + Args: + x: The input tenosr with ndim >= 2. + group_size: The group size used for quantization. + eps: The minimum to avoid dividing zero. + dtype: The dype of output tensor. Note that only `torch.int8` is supported for now. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. + """ + assert ( + x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.is_contiguous(), "`x` is not contiguous" + + iinfo = torch.iinfo(dtype) + int8_max = iinfo.max + int8_min = iinfo.min + + x_q = torch.empty_like(x, device=x.device, dtype=dtype) + M = x.numel() // group_size + N = group_size + x_s = torch.empty( + x.shape[:-1] + (x.shape[-1] // group_size,), + device=x.device, + dtype=torch.float32, + ) + + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + _per_token_group_quant_int8[(M,)]( + x, + x_q, + x_s, + group_size, + N, + eps, + int8_min=int8_min, + int8_max=int8_max, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=num_stages, + ) + + return x_q, x_s + + +def sglang_per_token_group_quant_int8( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: torch.dtype = torch.int8, +): + assert ( + x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.is_contiguous(), "`x` is not contiguous" + + iinfo = torch.iinfo(dtype) + int8_max = iinfo.max + int8_min = iinfo.min + + x_q = torch.empty_like(x, device=x.device, dtype=dtype) + x_s = torch.empty( + x.shape[:-1] + (x.shape[-1] // group_size,), + device=x.device, + dtype=torch.float32, + ) + + sgl_per_token_group_quant_int8(x, x_q, x_s, group_size, eps, int8_min, int8_max) + + return x_q, x_s + + +@triton.jit +def _w8a8_block_int8_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).to(tl.float32) * 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) + + +@functools.lru_cache +def get_w8a8_block_int8_configs( + N: int, K: int, block_n: int, block_k: int +) -> Optional[Dict[int, Any]]: + """ + Return optimized configurations for the w8a8 block fp8 kernel. + + The return value will be a dictionary that maps an irregular grid of + batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the + kernel on a given batch size bs, the closest batch size in the grid should + be picked and the associated configuration chosen to invoke the kernel. + """ + + # First look up if an optimized configuration is available in the configs + # directory + device_name = get_device_name().replace(" ", "_") + json_file_name = f"N={N},K={K},device_name={device_name},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json" + + config_file_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name + ) + if os.path.exists(config_file_path): + with open(config_file_path) as f: + logger.info( + "Using configuration from %s for W8A8 Block INT8 kernel.", + config_file_path, + ) + # If a configuration has been found, return it + return {int(key): val for key, val in json.load(f).items()} + + # If no optimized configuration is available, we will use the default + # configuration + logger.warning( + ( + "Using default W8A8 Block INT8 kernel config. Performance might be sub-optimal! " + "Config file not found at %s" + ), + config_file_path, + ) + return None + + +def w8a8_block_int8_matmul( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + """This function performs matrix multiplication with block-wise quantization. + + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + + Returns: + torch.Tensor: The result of matmul. + """ + assert len(block_size) == 2 + block_n, block_k = block_size[0], block_size[1] + + assert A.shape[-1] == B.shape[-1] + assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() + assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] + M = A.numel() // A.shape[-1] + + assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 + N, K = B.shape + assert triton.cdiv(N, block_n) == Bs.shape[0] + assert triton.cdiv(K, block_k) == Bs.shape[1] + + C_shape = A.shape[:-1] + (N,) + C = A.new_empty(C_shape, dtype=output_dtype) + + configs = get_w8a8_block_int8_configs(N, K, block_size[0], block_size[1]) + if configs: + # If an optimal configuration map has been found, look up the + # optimal config + config = configs[min(configs.keys(), key=lambda x: abs(x - M))] + else: + # Default config + # Block-wise quant: BLOCK_SIZE_K must be divisible by block_size[1] + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": block_size[0], + "BLOCK_SIZE_K": block_size[1], + "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_int8_matmul[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + **config, + ) + + return C diff --git a/python/sglang/srt/layers/quantization/int8_utils.py b/python/sglang/srt/layers/quantization/int8_utils.py new file mode 100644 index 000000000..9a7aa6d82 --- /dev/null +++ b/python/sglang/srt/layers/quantization/int8_utils.py @@ -0,0 +1,73 @@ +from typing import List, Optional, Tuple + +import torch + +from sglang.srt.layers.quantization.int8_kernel import ( + per_token_group_quant_int8, + w8a8_block_int8_matmul, +) + + +def apply_w8a8_block_int8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = per_token_group_quant_int8(input_2d, block_size[1]) + output = w8a8_block_int8_matmul( + q_input, weight, x_scale, weight_scale, block_size, output_dtype=input.dtype + ) + + if bias is not None: + output = output + bias + return output.to(dtype=input.dtype).view(*output_shape) + + +def input_to_int8( + x: torch.Tensor, dtype: torch.dtype = torch.int8 +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function quantizes input values to int8 values with tensor-wise quantization.""" + iinfo = torch.iinfo(dtype) + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) + int8_min, int8_max = iinfo.min, iinfo.max + scale = int8_max / amax + x_scl_sat = (x * scale).clamp(min=int8_min, max=int8_max) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + + +def block_dequant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: List[int], +) -> torch.Tensor: + """This function conducts block-wise dequantization. + The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale + and the block size. + The outputs are dequantized tensor. + """ + block_n, block_k = block_size[0], block_size[1] + n, k = x_q_block.shape + n_tiles = (n + block_n - 1) // block_n + k_tiles = (k + block_k - 1) // block_k + assert n_tiles == x_s.shape[0] + assert k_tiles == x_s.shape[1] + + x_dq_block = x_q_block.to(torch.float32) + + for i in range(k_tiles): + for j in range(n_tiles): + x_dq_block[ + j * block_n : min((j + 1) * block_n, n), + i * block_k : min((i + 1) * block_k, k), + ] *= x_s[j][i] + + return x_dq_block diff --git a/python/sglang/srt/layers/quantization/kv_cache.py b/python/sglang/srt/layers/quantization/kv_cache.py new file mode 100644 index 000000000..7da2dac17 --- /dev/null +++ b/python/sglang/srt/layers/quantization/kv_cache.py @@ -0,0 +1,82 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/kv_cache.py + +import logging + +import torch + +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.layers.radix_attention import RadixAttention + +logger = logging.getLogger(__name__) + + +class BaseKVCacheMethod(QuantizeMethodBase): + """ + Quant method that adds `k_scale` and `v_scale` attributes to the + Attention layer to support loading those scaling factors from checkpoints. + The k/v_scale will be used to: + - quantize k/v_cache entries before saving them to the cache + - dequantize k/v_cache entries before fetching them from the cache + + :param quant_config: the appropriate QuantizationConfig + """ + + def __init__(self, quant_config: QuantizationConfig): + self.quant_config = quant_config + + def create_weights(self, layer: torch.nn.Module): + """ + Create "weight" (aka k_scale and v_scale) for an attention layer. + """ + # Initialize the KV cache scales to -1.0, which is an invalid value. + # If the k/v_scale appears in the checkpoint, it will be + # overwritten when loading weights. + layer.k_scale = torch.nn.Parameter( + torch.tensor(-1.0, dtype=torch.float32), requires_grad=False + ) + layer.v_scale = torch.nn.Parameter( + torch.tensor(-1.0, dtype=torch.float32), requires_grad=False + ) + + def apply(self, layer: torch.nn.Module) -> torch.Tensor: + raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.") + + def process_weights_after_loading(self, layer: RadixAttention) -> None: + if layer.k_scale > 0.0 and layer.v_scale > 0.0: + # We prefer to use separate k_scale and v_scale if present + k_scale = layer.k_scale.to("cpu").tolist() + v_scale = layer.v_scale.to("cpu").tolist() + if is_fp8_fnuz(): + k_scale *= 2 + v_scale *= 2 + elif layer.k_scale < 0.0 and layer.v_scale < 0.0: + # If no scales were loaded (both scales are invalid negative + # values), use the default value of 1.0 + k_scale = 1.0 + v_scale = 1.0 + else: + # If we find a single kv_scale in the checkpoint, we remap + # kv_scale to k_scale during weight loading, and duplicate + # k_scale to v_scale here + assert layer.k_scale > 0.0 + scale_to_duplicate = max(layer.k_scale, layer.v_scale) + k_scale = scale_to_duplicate.to("cpu").tolist() + v_scale = scale_to_duplicate.to("cpu").tolist() + if is_fp8_fnuz(): + k_scale *= 2 + v_scale *= 2 + + if not isinstance(k_scale, float) or not isinstance(v_scale, float): + raise ValueError( + "Only support per-tensor scaling factor " "for fp8 KV cache" + ) + + # These are used in the final Attention.forward() + layer.k_scale.copy_(k_scale) + layer.v_scale.copy_(v_scale) + layer.k_scale_float = k_scale + layer.v_scale_float = v_scale diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py new file mode 100644 index 000000000..e0b398c25 --- /dev/null +++ b/python/sglang/srt/layers/quantization/marlin_utils.py @@ -0,0 +1,808 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/marlin_utils.py + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Optional + +import numpy +import torch + +from sglang.srt.layers.parameter import ( + BasevLLMParameter, + ChannelQuantScaleParameter, + GroupQuantScaleParameter, + PackedvLLMParameter, +) +from sglang.srt.layers.quantization.base_config import ( + LinearMethodBase, + QuantizationConfig, +) +from sglang.srt.layers.quantization.utils import ( + get_scalar_types, + pack_cols, + unpack_cols, +) +from sglang.srt.utils import get_device_capability, is_cuda + +if TYPE_CHECKING: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + +try: + from vllm import _custom_ops as ops +except ImportError: + ops = None + +_is_cuda = is_cuda() + +if _is_cuda: + from sgl_kernel import gptq_marlin_gemm + +logger = logging.getLogger(__name__) + +ScalarType, scalar_types = get_scalar_types() + +GPTQ_MARLIN_TILE = 16 +GPTQ_MARLIN_MIN_THREAD_N = 64 +GPTQ_MARLIN_MIN_THREAD_K = 128 +GPTQ_MARLIN_MAX_PARALLEL = 16 + +MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] + +# In case there is a performance issue with Marlin, the variable below can be +# changed to False, which allows Marlin to perform global reductions in fp16 +# precision (instead of fp32), and therefore, save on some memory movements. +USE_FP32_REDUCE_DEFAULT = True + + +# For binary size and compile time, we don't support the same types for with and +# without runtime zero-point. We support common cases, i.e. AWQ and GPTQ. +# TODO: we may want to move this into the C++ so its closer to the actual impl +def query_marlin_supported_quant_types( + has_zp: Optional[bool] = None, + include_fp_type: bool = True, + device_capability: Optional[int] = None, +): + if device_capability is None: + major, minor = get_device_capability() + capability = major * 10 + minor + device_capability = -1 if capability is None else capability + + if device_capability < 80: + return [] + + # - has_zp is True: return quant_types that has zero points + # - has_zp is False: return quant_types that has not zero points + # - has_zp is None: both + if has_zp is None: + types0 = query_marlin_supported_quant_types( + False, include_fp_type, device_capability + ) + types1 = query_marlin_supported_quant_types( + True, include_fp_type, device_capability + ) + return types0 + types1 + + if has_zp: + # AWQ style, unsigned + runtime zero-point + return [scalar_types.uint4] + else: + # GPTQ style, unsigned + symmetric bias + res = [scalar_types.uint4b8, scalar_types.uint8b128] + if include_fp_type: + res += [scalar_types.float8_e4m3fn, scalar_types.float4_e2m1f] + return res + + +def _check_marlin_supported( + quant_type: ScalarType, + group_size: Optional[int], + has_zp: bool, + device_capability: Optional[int] = None, +) -> tuple[bool, Optional[str]]: + + if device_capability is None: + major, minor = get_device_capability() + capability = major * 10 + minor + device_capability = -1 if capability is None else capability + + supported_types = query_marlin_supported_quant_types( + has_zp, True, device_capability + ) + + if quant_type not in supported_types: + return ( + False, + f"Marlin does not support weight_bits = {quant_type}. " + f"Only types = {supported_types} " + f"are supported (for group_size = {group_size}, " + f"device_capability = {device_capability}, zp = {has_zp}).", + ) + if group_size is None or group_size not in MARLIN_SUPPORTED_GROUP_SIZES: + return ( + False, + f"Marlin does not support group_size = {group_size}. " + f"Only group_sizes = {MARLIN_SUPPORTED_GROUP_SIZES} " + "are supported.", + ) + + return True, None + + +def check_marlin_supported( + quant_type: ScalarType, + group_size: int, + has_zp: bool = False, + device_capability: Optional[int] = None, +) -> bool: + cond, _ = _check_marlin_supported(quant_type, group_size, has_zp, device_capability) + return cond + + +def verify_marlin_supported( + quant_type: ScalarType, group_size: int, has_zp: bool = False +) -> None: + cond, err_msg = _check_marlin_supported(quant_type, group_size, has_zp) + if not cond: + assert err_msg is not None + raise ValueError(err_msg) + + +def verify_marlin_supports_shape( + output_size_per_partition: int, + input_size_per_partition: int, + input_size: int, + group_size: int, +) -> None: + + # Validate output_size_per_partition + if output_size_per_partition % GPTQ_MARLIN_MIN_THREAD_N != 0: + raise ValueError( + f"Weight output_size_per_partition = " + f"{output_size_per_partition} is not divisible by " + f" min_thread_n = {GPTQ_MARLIN_MIN_THREAD_N}. " + "Consider reducing tensor_parallel_size or running " + "with --quantization gptq." + ) + + # Validate input_size_per_partition + if input_size_per_partition % GPTQ_MARLIN_MIN_THREAD_K != 0: + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible " + f"by min_thread_k = {GPTQ_MARLIN_MIN_THREAD_K}. " + "Consider reducing tensor_parallel_size or running " + "with --quantization gptq." + ) + + if group_size < input_size and input_size_per_partition % group_size != 0: + raise ValueError( + f"Weight input_size_per_partition = {input_size_per_partition}" + f" is not divisible by group_size = {group_size}. " + "Consider reducing tensor_parallel_size or running " + "with --quantization gptq." + ) + + +def check_marlin_supports_shape( + output_size_per_partition: int, + input_size_per_partition: int, + input_size: int, + group_size: int, +) -> tuple[bool, Optional[str]]: + try: + verify_marlin_supports_shape( + output_size_per_partition, input_size_per_partition, input_size, group_size + ) + except ValueError as e: + return False, e.__str__() + return True, None + + +def check_marlin_supports_layer(layer: LinearBase, group_size: int) -> bool: + output_size_per_partition = ( + getattr(layer, "output_size_per_partition", None) or layer.output_size + ) + input_size_per_partition = ( + getattr(layer, "input_size_per_partition", None) or layer.input_size + ) + + return check_marlin_supports_shape( + output_size_per_partition=output_size_per_partition, + input_size_per_partition=input_size_per_partition, + input_size=layer.input_size, + group_size=group_size, + )[0] + + +def check_moe_marlin_supports_layer(layer: FusedMoE, group_size: int) -> bool: + hidden_size = layer.hidden_size + intermediate_size_per_partition = layer.intermediate_size_per_partition + # apply_router_weight_on_input is not supported for moe marlin + supports_router_weight = not layer.moe_runner_config.apply_router_weight_on_input + # moe marlin requires the activation to be silu + supports_activation = layer.moe_runner_config.activation == "silu" + + # gate-up: (n, k) = (intermediate_size_per_partition * 2, hidden_size) + # down: (n, k) = (hidden_size, intermediate_size_per_partition) + # moe marlin requires n % 128 == 0 and k % 64 == 0 + supports_shape = ( + hidden_size % 128 == 0 + and intermediate_size_per_partition % max(64, group_size) == 0 + ) + supports_group_size = group_size in [-1, 32, 64, 128] + return ( + supports_shape + and supports_group_size + and supports_router_weight + and supports_activation + ) + + +def marlin_make_workspace( + device: torch.device, max_blocks_per_sm: int = 1 +) -> torch.Tensor: + # In the new marlin kernel, we use the num of threadblocks as workspace + # size. The num of threadblocks is is sms_count * max_blocks_per_sm. + sms = torch.cuda.get_device_properties(device).multi_processor_count + return torch.zeros( + sms * max_blocks_per_sm, dtype=torch.int, device=device, requires_grad=False + ) + + +def marlin_is_k_full(act_order: bool, is_row_parallel: bool) -> bool: + return (not act_order) or (act_order and not is_row_parallel) + + +def marlin_repeat_scales_on_all_ranks( + act_order: bool, group_size: int, is_row_parallel: bool +) -> bool: + # Need to repeat scales on every rank if act_ordering or + # channelwise and RowParallelLinear + is_channelwise = group_size == -1 + return act_order or (is_channelwise and is_row_parallel) + + +def marlin_make_empty_g_idx(device: torch.device) -> torch.Tensor: + return torch.nn.Parameter( + torch.empty(0, dtype=torch.int, device=device), requires_grad=False + ) + + +def marlin_make_empty_zp(device: torch.device) -> torch.Tensor: + return torch.nn.Parameter( + torch.empty(0, dtype=torch.int, device=device), requires_grad=False + ) + + +def marlin_sort_g_idx(g_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + g_idx_sort_indices = torch.argsort(g_idx).to(torch.int) + return g_idx[g_idx_sort_indices], g_idx_sort_indices + + +def get_scale_perms(): + scale_perm: list[int] = [] + for i in range(8): + scale_perm.extend([i + 8 * j for j in range(8)]) + scale_perm_single: list[int] = [] + for i in range(4): + scale_perm_single.extend([2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]]) + return scale_perm, scale_perm_single + + +def marlin_permute_scales( + s: torch.Tensor, size_k: int, size_n: int, group_size: int +) -> torch.Tensor: + + scale_perm, scale_perm_single = get_scale_perms() + if group_size < size_k and group_size != -1: + s = s.reshape((-1, len(scale_perm)))[:, scale_perm] + else: + s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single] + s = s.reshape((-1, size_n)).contiguous() + + return s + + +def marlin_permute_bias(s: torch.Tensor) -> torch.Tensor: + origin_shape = s.shape + _, scale_perm_single = get_scale_perms() + s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single] + return s.reshape(*origin_shape).contiguous() + + +def marlin_moe_permute_scales( + s: torch.Tensor, + size_k: int, + size_n: int, + group_size: int, +): + num_experts = s.shape[0] + output = torch.empty( + (num_experts, s.shape[1], s.shape[2]), + device=s.device, + dtype=s.dtype, + ) + + for e in range(num_experts): + output[e] = marlin_permute_scales(s[e], size_k, size_n, group_size) + return output + + +def marlin_zero_points( + zp: torch.Tensor, size_k: int, size_n: int, num_bits: int +) -> torch.Tensor: + # Permute zero-points in a similar way to scales, but do not use the + # "single" permutation, since zero-points are applied on every MMA + scale_perm, _ = get_scale_perms() + zp = zp.reshape((-1, len(scale_perm)))[:, scale_perm] + + # Interleave column dim (for the dequantize code) and pack it to int32 + if num_bits == 4: + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = numpy.array([0, 2, 1, 3]) + else: + raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) + + zp = zp.reshape((-1, len(interleave)))[:, interleave].ravel() + zp = zp.reshape((-1, size_n)).contiguous() + zp = pack_cols(zp, num_bits, size_k, size_n) + + return zp + + +def awq_to_marlin_zero_points( + q_zp_packed: torch.Tensor, size_k: int, size_n: int, num_bits: int +) -> torch.Tensor: + # AWQ zero-points are quantized and packed on the column dim. + # In addition, the values are permuted based on dequantizer. + # Here we undo both of these, and then apply marlin permutation + # and pack it back. + q_zp = unpack_cols(q_zp_packed, num_bits, size_k, size_n) + + # Undo interleaving (use argsort(..) to get inverse perm) + if num_bits == 4: + undo_interleave = numpy.argsort(numpy.array([0, 2, 4, 6, 1, 3, 5, 7])) + elif num_bits == 8: + undo_interleave = numpy.argsort(numpy.array([0, 2, 1, 3])) + else: + raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) + + q_zp = q_zp.reshape((-1, len(undo_interleave)))[:, undo_interleave].ravel() + q_zp = q_zp.reshape((-1, size_n)).contiguous() + + marlin_zp = marlin_zero_points(q_zp, size_k, size_n, num_bits) + return marlin_zp + + +def moe_awq_to_marlin_zero_points( + q_zp_packed: torch.Tensor, size_k: int, size_n: int, num_bits: int +): + num_experts = q_zp_packed.shape[0] + output = torch.empty( + (num_experts, q_zp_packed.shape[1], q_zp_packed.shape[2]), + device=q_zp_packed.device, + dtype=q_zp_packed.dtype, + ) + for e in range(num_experts): + output[e] = awq_to_marlin_zero_points(q_zp_packed[e], size_k, size_n, num_bits) + return output + + +def maybe_warn_marlin_atomic_add(device, dtype): + if torch.compiler.is_dynamo_compiling(): + return + device_capability = torch.cuda.get_device_capability(device) + if device_capability[0] < 9 and dtype == torch.bfloat16: + logger.info_once( + "You are running Marlin kernel with bf16 on GPUs before SM90. " + "You can consider change to fp16 to achieve better performance " + "if possible." + ) + + +def maybe_warn_marlin_atomic_add_env(): + if torch.compiler.is_dynamo_compiling(): + return + # TODO(yiyun): Need to add sglang's MARLIN_USE_ATOMIC_ADD: bool = False + if True: + return + # if envs.VLLM_MARLIN_USE_ATOMIC_ADD: + # return + logger.info_once( + "Marlin kernel can achieve better performance for small size_n " + "with experimental use_atomic_add feature. " + "You can consider set environment variable " + "VLLM_MARLIN_USE_ATOMIC_ADD to 1 if possible." + ) + + +def should_use_atomic_add_reduce( + m: int, n: int, k: int, device: torch.device, dtype: torch.dtype +) -> bool: + + # the performance of atomicAdd is better than global reduce + # only when m*n is small and k is large + if n >= 2048 or k < 2048 or device.type != "cuda": + return False + + # disable atomicAdd reduce by default, + # one can enable it with VLLM_MARLIN_USE_ATOMIC_ADD=1 + # TODO: Need to add sglang's MARLIN_USE_ATOMIC_ADD: bool = False + if not True: + maybe_warn_marlin_atomic_add_env() + return False + + # sm8x doesn't support atomicAdd + bfloat16 natively + device_capability = torch.cuda.get_device_capability(device) + if device_capability[0] < 9 and dtype == torch.bfloat16: + maybe_warn_marlin_atomic_add(device, dtype) + return False + + return True + + +def apply_gptq_marlin_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_zp: torch.Tensor, + g_idx: torch.Tensor, + g_idx_sort_indices: torch.Tensor, + workspace: torch.Tensor, + wtype: ScalarType, + output_size_per_partition: int, + input_size_per_partition: int, + is_k_full: bool, + bias: Optional[torch.Tensor] = None, + use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, +) -> torch.Tensor: + reshaped_x = input.reshape(-1, input.shape[-1]) + out_shape = input.shape[:-1] + (output_size_per_partition,) + + use_atomic_add = should_use_atomic_add_reduce( + m=reshaped_x.size(0), + n=output_size_per_partition, + k=reshaped_x.size(1), + device=input.device, + dtype=input.dtype, + ) + + output = gptq_marlin_gemm( + reshaped_x, + None, + weight, + weight_scale, + None, + weight_zp, + g_idx, + g_idx_sort_indices, + workspace, + wtype, + size_m=reshaped_x.shape[0], + size_n=output_size_per_partition, + size_k=input_size_per_partition, + is_k_full=is_k_full, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + is_zp_float=False, + ) + + if bias is not None: + output.add_(bias) # In-place add + + return output.reshape(out_shape) + + +def apply_awq_marlin_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_zp: torch.Tensor, + g_idx: torch.Tensor, + g_idx_sort_indices: torch.Tensor, + workspace: torch.Tensor, + quant_type: ScalarType, + output_size_per_partition: int, + input_size_per_partition: int, + bias: Optional[torch.Tensor] = None, + use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, +) -> torch.Tensor: + reshaped_x = input.reshape(-1, input.shape[-1]) + out_shape = input.shape[:-1] + (output_size_per_partition,) + + use_atomic_add = should_use_atomic_add_reduce( + m=reshaped_x.size(0), + n=output_size_per_partition, + k=reshaped_x.size(1), + device=input.device, + dtype=input.dtype, + ) + + output = gptq_marlin_gemm( + reshaped_x, + None, + weight, + weight_scale, + None, + weight_zp, + g_idx, + g_idx_sort_indices, + workspace, + quant_type, + size_m=reshaped_x.shape[0], + size_n=output_size_per_partition, + size_k=input_size_per_partition, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + is_zp_float=False, + ) + + if bias is not None: + output.add_(bias) # In-place add + + return output.reshape(out_shape) + + +class MarlinConfig(QuantizationConfig): + """Config class for Marlin. + + Reference: https://github.com/IST-DASLab/marlin/tree/master + """ + + def __init__( + self, + group_size: int, + lm_head_quantized: bool, + ) -> None: + super().__init__() + + # Group size for the quantization. + self.group_size = group_size + self.lm_head_quantized = lm_head_quantized + if self.group_size != 128 and self.group_size != -1: + raise ValueError( + "Currently, only group size 128 and -1 (channelwise) " + "is supported for Marlin, but got group_size of " + f"{self.group_size}" + ) + + # 4 Bits packed into 32 bit datatype. + self.pack_factor = 32 // 4 + + # Tile size used by marlin kernels. + self.tile_size = 16 + + # Min out_features dim + self.min_n_threads = 64 + + # Min in_features dim + self.min_k_threads = 128 + + # Max parallel problems to solve at once (improves large + # batch performance) + self.max_parallel = 16 + + # Permutation length used by the marlin kernels. + self.perm_len = 1024 + + def __repr__(self) -> str: + return ( + f"MarlinConfig(group_size={self.group_size}, " + f"lm_head_quantized={self.lm_head_quantized})" + ) + + @classmethod + def get_name(cls) -> str: + return "marlin" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.half] + + @classmethod + # Need to figure it out + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return ["quantize_config.json"] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "MarlinConfig": + group_size = cls.get_from_keys(config, ["group_size"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + return cls(group_size, lm_head_quantized) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + # compat: autogptq >=0.8.0 use checkpoint_format: str + # compat: autogptq <=0.7.1 is_marlin_format: bool + is_marlin_format = hf_quant_cfg.get( + "checkpoint_format" + ) == "marlin" or hf_quant_cfg.get("is_marlin_format", False) + + is_valid_user_quant = ( + user_quant is None or user_quant == "gptq" or user_quant == "marlin" + ) + + if is_marlin_format and is_valid_user_quant: + msg = "The model is serialized in {} format. Using {} kernel.".format( + cls.get_name(), cls.get_name() + ) + logger.info(msg) + return cls.get_name() + + return None + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[MarlinLinearMethod]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + + if isinstance(layer, LinearBase) or ( + isinstance(layer, ParallelLMHead) and self.lm_head_quantized + ): + return MarlinLinearMethod(self) + return None + + +class MarlinLinearMethod(LinearMethodBase): + """Linear method for Marlin. + + Args: + quant_config: The Marlin quantization config. + """ + + def __init__(self, quant_config: MarlinConfig): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + del output_size # Unused. + weight_loader = extra_weight_attrs["weight_loader"] + + if params_dtype != torch.float16: + raise ValueError( + f"The params dtype must be float16, but got {params_dtype}" + ) + + # Validate output_size_per_partition + output_size_per_partition = sum(output_partition_sizes) + if output_size_per_partition % self.quant_config.min_n_threads != 0: + raise ValueError( + f"Weight output_size_per_partition = " + f"{output_size_per_partition} is not divisible by " + f"min_n_threads = {self.quant_config.min_n_threads}." + ) + if output_size_per_partition % self.quant_config.pack_factor != 0: + raise ValueError( + f"Weight output_size_per_partition = " + f"{output_size_per_partition} is not divisible by " + f"pack_factor = {self.quant_config.pack_factor}." + ) + + # Validate input_size_per_partition + if input_size_per_partition % self.quant_config.min_k_threads != 0: + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"min_k_threads = {self.quant_config.min_k_threads}." + ) + if ( + self.quant_config.group_size != -1 + and input_size_per_partition % self.quant_config.group_size != 0 + ): + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"group_size = {self.quant_config.group_size}." + ) + + # Check that we have at least 4 tiles horizontally in the shard + num_tiles_per_perm = self.quant_config.perm_len // ( + self.quant_config.tile_size**2 + ) + if output_size_per_partition % num_tiles_per_perm != 0: + raise ValueError("Each permutation group must reside on the same gpu") + + # Quantized 4Bit weights packed into Int32. + qweight = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition // self.quant_config.tile_size, + output_size_per_partition + * self.quant_config.tile_size + // self.quant_config.pack_factor, + device="cuda", + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + marlin_tile_size=self.quant_config.tile_size, + weight_loader=weight_loader, + ) + + # Determine if channelwise or not + input_groups = ( + 1 + if self.quant_config.group_size == -1 + else input_size_per_partition // self.quant_config.group_size + ) + + weight_scale_args = { + "data": torch.empty( + input_groups, + output_size_per_partition, + device="cuda", + dtype=params_dtype, + ), + "weight_loader": weight_loader, + } + if input_groups == 1: + scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args) + else: + scales = GroupQuantScaleParameter( + output_dim=1, input_dim=0, **weight_scale_args + ) + + # Allocate workspace (Used for internal locking mechanism) + max_workspace_size = ( + output_size_per_partition // self.quant_config.min_n_threads + ) * self.quant_config.max_parallel + + workspace = BasevLLMParameter( + data=torch.zeros(max_workspace_size, device="cuda", dtype=torch.int), + weight_loader=weight_loader, + ) + + layer.register_parameter("B", qweight) + layer.register_parameter("s", scales) + layer.register_parameter("workspace", workspace) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # required by torch.compile + layer.B = torch.nn.Parameter(layer.B.data, requires_grad=False) + layer.s = torch.nn.Parameter(layer.s.data, requires_grad=False) + layer.workspace = torch.nn.Parameter(layer.workspace.data, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + qweight = layer.B + scales = layer.s + workspace = layer.workspace + + x_2d = x.view(-1, x.shape[-1]) + + size_m = x_2d.shape[0] + size_k = x_2d.shape[1] + size_n = scales.shape[1] + + output_2d = ops.marlin_gemm( + x_2d, qweight, scales, workspace, size_m, size_n, size_k + ) + + output = output_2d.view(x.shape[:-1] + (output_2d.shape[1],)) + + if bias is not None: + output.add_(bias) # In-place add + + return output diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp8.py b/python/sglang/srt/layers/quantization/marlin_utils_fp8.py new file mode 100644 index 000000000..94326d71e --- /dev/null +++ b/python/sglang/srt/layers/quantization/marlin_utils_fp8.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: Apache-2.0 + +import logging +from typing import Optional + +import torch + +from sglang.srt.layers.quantization.marlin_utils import ( + USE_FP32_REDUCE_DEFAULT, + marlin_make_workspace, + marlin_permute_bias, + marlin_permute_scales, + should_use_atomic_add_reduce, +) +from sglang.srt.layers.quantization.utils import get_scalar_types +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() +if _is_cuda: + from sgl_kernel import gptq_marlin_gemm, gptq_marlin_repack + +ScalarType, scalar_types = get_scalar_types() + +logger = logging.getLogger(__name__) + + +def fp8_fused_exponent_bias_into_scales(scales): + fp8_exponent = 4 + if scales.dtype == torch.half: + target_exponent = 5 + elif scales.dtype == torch.bfloat16: + target_exponent = 8 + # exponent_bias_fp16 = 2 ** 4 - 2 ** 3 = 8 + # exponent_bias_bf16 = 2 ** 7 - 2 ** 3 = 120 + exponent_bias = 2 ** (target_exponent - 1) - 2 ** (fp8_exponent - 1) + s = torch.ones_like(scales) * 2 + s = s**exponent_bias + return scales * s + + +def apply_fp8_marlin_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + workspace: torch.Tensor, + size_n: int, + size_k: int, + bias: Optional[torch.Tensor], + use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, +) -> torch.Tensor: + # For GPUs that lack FP8 hardware support, we can leverage the + # Marlin kernel for fast weight-only FP8 quantization + + reshaped_x = input.reshape(-1, input.shape[-1]) + out_shape = input.shape[:-1] + (size_n,) + + use_atomic_add = should_use_atomic_add_reduce( + m=reshaped_x.size(0), n=size_n, k=size_k, device=input.device, dtype=input.dtype + ) + + output = gptq_marlin_gemm( + a=reshaped_x, + c=None, + b_q_weight=weight, + b_bias=bias, + b_scales=weight_scale, + global_scale=None, + b_zeros=None, + g_idx=None, + perm=None, + workspace=workspace, + b_q_type=scalar_types.float8_e4m3fn, + size_m=reshaped_x.size(0), + size_n=size_n, + size_k=size_k, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + ) + + return output.reshape(out_shape) + + +def prepare_fp8_layer_for_marlin( + layer: torch.nn.Module, size_k_first: bool = True +) -> None: + logger.warning_once( + "Your GPU does not have native support for FP8 computation but " + "FP8 quantization is being used. Weight-only FP8 compression will " + "be used leveraging the Marlin kernel. This may degrade " + "performance for compute-heavy workloads." + ) + + part_size_n = layer.output_size_per_partition + part_size_k = layer.input_size_per_partition + weight_block_size = getattr(layer, "weight_block_size", None) + + if size_k_first: + assert layer.weight.shape == (part_size_k, part_size_n) + else: + assert layer.weight.shape == (part_size_n, part_size_k) + + device = layer.weight.device + + # WORKSPACE + layer.workspace = marlin_make_workspace(device) + + # WEIGHT + # Repack weights to marlin format + perm = torch.empty(0, dtype=torch.int, device=device) + qweight = pack_fp8_to_int32(layer.weight, size_k_first) + if not size_k_first: + qweight = qweight.T.contiguous() + + marlin_qweight = gptq_marlin_repack( + b_q_weight=qweight, + perm=perm, + size_k=part_size_k, + size_n=part_size_n, + num_bits=8, + ) + layer.weight = torch.nn.Parameter(marlin_qweight, requires_grad=False) + + # WEIGHT SCALES + # Permute scales + if "weight_scale" in dir(layer): + scales = layer.weight_scale.to(layer.orig_dtype) + elif "weight_scale_inv" in dir(layer): + scales = layer.weight_scale_inv.to(layer.orig_dtype) + del layer.weight_scale_inv + + group_size = -1 if weight_block_size is None else weight_block_size[1] + + # marlin kernel only support channel-wise and group-wise quantization + # we need to convert the scales + if weight_block_size is None: + if scales.nelement() == 1: + # tensor-wise quantization -> channel-wise quantization + # (1, 1) =>(repeat)=> (1, size_n) + scales = scales.view(1, 1).repeat_interleave(part_size_n, 1) + elif scales.nelement() > 1 and scales.nelement() != part_size_n: + assert part_size_n % scales.nelement() == 0 + s_size = scales.nelement() + # tensor-wise quantization (for gate-up proj) + # -> channel-wise quantization + # (1, s_size) =>(repeat)=> (1, size_n) + scales = scales.view(1, s_size) + scales = scales.repeat_interleave(part_size_n // s_size, 1) + else: + # channel-wise quantization + # (1, size_n) + scales = scales.view(1, part_size_n) + else: + # block-wise quantization -> group-wise quantization + # (size_k // block_size[1], ceil(size_n / block_size[0])) + # =>(repeat)=> (size_k // block_size[1], size_n) + if not size_k_first: + scales = scales.T.contiguous() + block_n = weight_block_size[0] + scales = scales.repeat_interleave(block_n, 1) + # size_n may not divisible by block_size[0] + scales = scales[:, :part_size_n] + + marlin_scales = marlin_permute_scales( + s=scales, size_k=part_size_k, size_n=part_size_n, group_size=group_size + ) + marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales) + layer.weight_scale = torch.nn.Parameter(marlin_scales, requires_grad=False) + + if hasattr(layer, "bias") and layer.bias is not None: + assert layer.bias.shape == (part_size_n,) + bias = marlin_permute_bias(layer.bias) + layer.bias = torch.nn.Parameter(bias, requires_grad=False) + + +def prepare_moe_fp8_layer_for_marlin( + layer: torch.nn.Module, size_k_first: bool = True +) -> None: + logger.warning_once( + "Your GPU does not have native support for FP8 computation but " + "FP8 quantization is being used. Weight-only FP8 compression will " + "be used leveraging the Marlin kernel. This may degrade " + "performance for compute-heavy workloads." + ) + + e = layer.num_experts + k = layer.hidden_size + n = layer.intermediate_size_per_partition + weight_block_size = getattr(layer, "weight_block_size", None) + + # WORKSPACE + device = layer.w13_weight.device + layer.workspace = marlin_make_workspace(device, 4) + perm = torch.empty(0, dtype=torch.int, device=device) + + # WEIGHT + # Repack weights to marlin format + for name in ["w13_weight", "w2_weight"]: + weight = getattr(layer, name) + tensor_list = [] + if "w13" in name: + size_n, size_k = n * 2, k + else: + size_n, size_k = k, n + + if size_k_first: + assert weight.shape == (e, size_k, size_n) + else: + assert weight.shape == (e, size_n, size_k) + + for i in range(e): + qweight = pack_fp8_to_int32(weight[i], size_k_first) + if not size_k_first: + qweight = qweight.T.contiguous() + + marlin_qweight = gptq_marlin_repack( + b_q_weight=qweight, perm=perm, size_k=size_k, size_n=size_n, num_bits=8 + ) + tensor_list.append(marlin_qweight) + + weight = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + weight = torch.nn.Parameter(weight, requires_grad=False) + + setattr(layer, name, weight) + + # WEIGHT SCALES + # Permute scales + group_size = -1 if weight_block_size is None else weight_block_size[1] + + for name in ["w13", "w2"]: + if name + "_weight_scale" in dir(layer): + new_name = name + "_weight_scale" + scales = getattr(layer, new_name).to(layer.orig_dtype) + delattr(layer, new_name) + elif name + "_weight_scale_inv" in dir(layer): + new_name = name + "_weight_scale_inv" + scales = getattr(layer, new_name).to(layer.orig_dtype) + delattr(layer, new_name) + + tensor_list = [] + if "w13" in name: + size_n, size_k = n * 2, k + else: + size_n, size_k = k, n + + # marlin kernel only support channel-wise and group-wise quantization + # we need to convert the scales + if weight_block_size is None: + if scales.nelement() == e: + # tensor-wise quantization -> channel-wise quantization + # (e, 1, 1) =>(repeat)=> (e, 1, size_n) + scales = scales.view(e, 1, 1).repeat_interleave(size_n, 2) + elif scales.nelement() > e and scales.nelement() != e * size_n: + assert (e * size_n) % scales.nelement() == 0 + s_size = scales.nelement() // e + # tensor-wise quantization (for gate-up proj) + # -> channel-wise quantization + # (e, 1, s_size) =>(repeat)=> (e, 1, size_n) + scales = scales.view(e, 1, s_size) + scales = scales.repeat_interleave(size_n // s_size, 2) + else: + # channel-wise quantization + # (e, 1, size_n) + scales = scales.view(e, 1, size_n) + else: + # block-wise quantization -> group-wise quantization + # (e, size_k // block_size[1], ceil(size_n / block_size[0])) + # =>(repeat)=> (e, size_k // block_size[1], size_n) + if not size_k_first: + scales = scales.permute(0, 2, 1) + block_n = weight_block_size[0] + scales = scales.repeat_interleave(block_n, 2) + # size_n may not divisible by block_size[0] + scales = scales[..., :size_n].contiguous() + + for i in range(e): + marlin_scales = marlin_permute_scales( + s=scales[i], size_k=size_k, size_n=size_n, group_size=group_size + ) + tensor_list.append(marlin_scales) + + scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + scales = fp8_fused_exponent_bias_into_scales(scales) + scales = torch.nn.Parameter(scales, requires_grad=False) + + setattr(layer, name + "_weight_scale", scales) + + # BIAS + # Permute bias + for name in ["w13_bias", "w2_bias"]: + if not hasattr(layer, name): + continue + bias = getattr(layer, name).to(layer.orig_dtype) + + tensor_list = [] + for i in range(e): + expert_bias = bias[i] + + tensor_list.append(marlin_permute_bias(expert_bias)) + + bias = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + bias = torch.nn.Parameter(bias, requires_grad=False) + setattr(layer, name, bias) + + +def pack_fp8_to_int32( + fp8_tensor: torch.Tensor, size_k_first: bool = True +) -> torch.Tensor: + """ + Repack FP8 weights to gptq format (packed int32 elements) + """ + assert fp8_tensor.dtype == torch.float8_e4m3fn + assert fp8_tensor.ndim == 2 + + fp8_tensor = fp8_tensor.T if size_k_first else fp8_tensor + fp8_tensor = fp8_tensor.contiguous() + # fp8_tensor is contiguous and have shape (N, K) now + # with `.view(torch.int32)`, it become (N, K // 4) + int32_tensor = fp8_tensor.view(torch.int32) + return int32_tensor.T.contiguous() if size_k_first else int32_tensor + + +def marlin_quant_fp8_torch(weight, group_size): + size_n, size_k = weight.shape + device = weight.device + + if group_size != -1: + scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 448 + repeated_scales = scales.repeat_interleave(group_size, 1) + fp8_weight = (weight / repeated_scales).to(torch.float8_e4m3fn) + weight_ref = fp8_weight.to(weight.dtype) * repeated_scales + else: + scales = weight.view(size_n, 1, group_size).abs().max(-1)[0] / 448 + repeated_scales = scales.repeat_interleave(size_k, 1) + fp8_weight = (weight / repeated_scales).to(torch.float8_e4m3fn) + weight_ref = fp8_weight.to(weight.dtype) * repeated_scales + + packed_weight = pack_fp8_to_int32(fp8_weight, False).T.contiguous() + marlin_qweight = gptq_marlin_repack( + b_q_weight=packed_weight, + perm=torch.empty(0, dtype=torch.int, device=device), + size_k=size_k, + size_n=size_n, + num_bits=8, + ) + + marlin_scales = marlin_permute_scales( + s=scales.T, size_k=size_k, size_n=size_n, group_size=group_size + ) + + marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales) + + return weight_ref.T, marlin_qweight, marlin_scales diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py new file mode 100755 index 000000000..89ecd44f5 --- /dev/null +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -0,0 +1,1402 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import torch +from torch.nn.parameter import Parameter + +from sglang.srt.distributed import get_tp_group +from sglang.srt.layers.dp_attention import get_dp_global_num_tokens, get_local_dp_buffer +from sglang.srt.layers.moe import ( + MoeRunner, + MoeRunnerBackend, + MoeRunnerConfig, + should_use_flashinfer_cutlass_moe_fp4_allgather, + should_use_flashinfer_trtllm_moe, +) +from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.parameter import ModelWeightParameter, PerTensorScaleParameter +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + cutlass_fp8_supported, + is_sm100_supported, +) +from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.layers.quantization.utils import ( + convert_to_channelwise, + is_layer_skipped, + per_tensor_dequantize, + requantize_with_max_scale, +) +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.utils import is_cuda, next_power_of_2 + +if TYPE_CHECKING: + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +if is_cuda(): + from sgl_kernel import scaled_fp4_quant + +try: + from flashinfer import mm_fp4 as fp4_gemm + from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_sf_a + + enable_flashinfer_fp4_gemm = True +except ImportError: + if is_cuda(): + from sgl_kernel import cutlass_scaled_fp4_mm as fp4_gemm + else: + fp4_gemm = None + enable_flashinfer_fp4_gemm = False + reorder_rows_for_gated_act_gemm = None + shuffle_matrix_a = None + shuffle_matrix_sf_a = None + +try: + from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe +except ImportError: + flashinfer_cutlass_fused_moe = None + +# Initialize logger for the module +logger = logging.getLogger(__name__) + +# Supported activation schemes for the current configuration +ACTIVATION_SCHEMES = ["static"] + + +class ModelOptFp8Config(QuantizationConfig): + """Configuration for ModelOpt FP8 quantization, including serialization and compatibility checks.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = False, + kv_cache_quant_method: Optional[str] = None, + exclude_modules: Optional[List[str]] = None, + ) -> None: + """ + Args: + is_checkpoint_fp8_serialized (bool): Indicates if the checkpoint uses serialized FP8 format. + """ + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + self.kv_cache_quant_method = kv_cache_quant_method + self.exclude_modules = exclude_modules + if is_checkpoint_fp8_serialized: + logger.warning( + "Detected ModelOpt FP8 checkpoint. The format is experimental and subject to change." + ) + + @classmethod + def get_name(cls) -> str: + return "modelopt" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 89 # Minimum hardware capability (e.g., Hopper GPUs). + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["hf_quant_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> ModelOptFp8Config: + # Handle two different config formats: + # 1. hf_quant_config.json format: {"quantization": {"quant_algo": "FP8", ...}} + # 2. config.json quantization_config format: {"quant_algo": "FP8", ...} + # In future modelopt will deprecate hf_quant_config.json, and only keep config.json. + # For legacy reasons, we keep hf_quant_config.json for now. + + # Initialize variables + kv_cache_quant_method = None + exclude_modules = None + + # Try flat format first (config.json quantization_config - preferred format) + quant_method = config.get("quant_algo") + if quant_method is not None: + # Flat format (config.json quantization_config) + # For kv_cache, check if kv_cache_scheme exists and extract algo + kv_cache_scheme = config.get("kv_cache_scheme") + if ( + kv_cache_scheme + and kv_cache_scheme.get("type") == "float" + and kv_cache_scheme.get("num_bits") == 8 + ): + kv_cache_quant_method = "FP8" + + # Map 'ignore' field to 'exclude_modules' + exclude_modules = config.get("ignore") + else: + # Fall back to nested format (hf_quant_config.json - legacy format) + try: + quantization_section = cls.get_from_keys(config, ["quantization"]) + quant_method = quantization_section.get("quant_algo") + kv_cache_quant_method = quantization_section.get("kv_cache_quant_algo") + exclude_modules = quantization_section.get("exclude_modules") + except ValueError: + raise ValueError( + "Cannot find 'quant_algo' in the model's quantization config. " + "Expected either flat format (config.json) or nested format (hf_quant_config.json)." + ) + if quant_method is None: + raise ValueError( + "Cannot find 'quant_algo' in the model's quantization config. " + ) + if "FP8" not in quant_method: + raise ValueError( + "ModelOptFp8Config only supports static FP8 quantization in SGLang. " + "For FP4 quantization, use ModelOptFp4Config. " + "Check the quantization config for your model's configuration." + ) + + return cls( + is_checkpoint_fp8_serialized=True, + kv_cache_quant_method=kv_cache_quant_method, + exclude_modules=exclude_modules, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if self.exclude_modules and any( + module in prefix + or ( + prefix.startswith("language_model.") + and module in prefix.removeprefix("language_model.") + ) + for module in self.exclude_modules + ): + return None + + if isinstance(layer, LinearBase): + return ModelOptFp8LinearMethod(self) + if self.kv_cache_quant_method and isinstance(layer, RadixAttention): + return ModelOptFp8KVCacheMethod(self) + + if isinstance(layer, FusedMoE): + return ModelOptFp8MoEMethod(self) + + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class ModelOptFp8LinearMethod(LinearMethodBase): + """Linear method for ModelOpt static FP8 quantization. + + Supports loading FP8 checkpoints with static weight and activation scales. + Future support may include dynamic scales. + + **Limitations**: + 1. Only supports per-tensor quantization due to `torch._scaled_mm` limitations. + 2. Only supports the `float8_e4m3fn` data type. + + Args: + quant_config (ModelOptFp8Config): The ModelOpt quantization configuration. + """ + + def __init__(self, quant_config: ModelOptFp8Config): + super().__init__() + self.quant_config = quant_config + self.cutlass_fp8_supported = cutlass_fp8_supported() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + """Creates and registers weights, weight scales, and input scales for FP8 quantization.""" + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_fp8_serialized + else params_dtype + ) + + # Set layer attributes + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + + # Register weight + layer.register_parameter( + "weight", + ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + if self.quant_config.is_checkpoint_fp8_serialized: + # Register weight and input scales + for scale_name in ["weight_scale", "input_scale"]: + layer.register_parameter( + scale_name, + PerTensorScaleParameter( + data=torch.full( + (len(output_partition_sizes),), + torch.finfo(torch.float32).min, + dtype=torch.float32, + ), + weight_loader=weight_loader, + ), + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Requantizes weights after loading using the maximum scale.""" + max_w_scale, quantized_weight = requantize_with_max_scale( + layer.weight, layer.weight_scale, layer.logical_widths + ) + layer.weight = Parameter(quantized_weight.t(), requires_grad=False) + # cutlass sgl-kernel only supports per-channel scale + if self.cutlass_fp8_supported: + max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths) + layer.weight_scale = Parameter(max_w_scale, requires_grad=False) + layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Applies FP8 linear transformation.""" + return apply_fp8_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + ) + + +class ModelOptFp8KVCacheMethod(BaseKVCacheMethod): + """ + Handles loading FP8 kv-cache scaling factors from modelopt quantized checkpoints. + """ + + def __init__(self, quant_config: ModelOptFp8Config): + super().__init__(quant_config) + + +class ModelOptFp8MoEMethod(FusedMoEMethodBase): + """MoE method for ModelOpt FP8. + Supports loading FP8 checkpoints with static weight scale and activation scale. + + Args: + quant_config: The ModelOpt quantization config. + """ + + def __init__(self, quant_config: ModelOptFp8Config): + self.quant_config = quant_config + self.cutlass_fp8_supported = cutlass_fp8_supported() + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + # Use FP8 dtype if checkpoint is serialized, otherwise use the default dtype + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_fp8_serialized + else params_dtype + ) + weight_loader = extra_weight_attrs.get("weight_loader") + + w13_weight = ModelWeightParameter( + data=torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=weight_dtype, + ), + input_dim=2, + output_dim=1, + weight_loader=weight_loader, + ) + layer.register_parameter("w13_weight", w13_weight) + + w2_weight = ModelWeightParameter( + data=torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=weight_dtype, + ), + input_dim=2, + output_dim=1, + weight_loader=weight_loader, + ) + layer.register_parameter("w2_weight", w2_weight) + + if self.quant_config.is_checkpoint_fp8_serialized: + # WEIGHT SCALES - Per-tensor scaling for ModelOpts + # Allocate 2 scales for w1 and w3 respectively. + # They will be combined to a single scale after weight loading. + w13_weight_scale = PerTensorScaleParameter( + data=torch.full( + (num_experts, 2), + torch.finfo(torch.float32).min, + dtype=torch.float32, + ), + weight_loader=weight_loader, + ) + w2_weight_scale = PerTensorScaleParameter( + data=torch.full( + (num_experts,), torch.finfo(torch.float32).min, dtype=torch.float32 + ), + weight_loader=weight_loader, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + # Set weight loader attributes for scales + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + + # INPUT SCALES - Per-tensor scaling for ModelOpt + w13_input_scale = PerTensorScaleParameter( + data=torch.full((num_experts,), 1.0, dtype=torch.float32), + weight_loader=weight_loader, + ) + w2_input_scale = PerTensorScaleParameter( + data=torch.full((num_experts,), 1.0, dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + layer.register_parameter("w2_input_scale", w2_input_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Process FP8 MoE weights after loading from serialized checkpoint. + + Only supports pre-quantized checkpoints with FP8 weights and scales. + """ + + layer.w13_weight = Parameter(layer.w13_weight.data, requires_grad=False) + layer.w2_weight = Parameter(layer.w2_weight.data, requires_grad=False) + + # Handle scale parameters + if hasattr(layer, "w13_weight_scale") and layer.w13_weight_scale is not None: + # Fp8 moe kernel needs single weight scale for w13 per expert. + # We take the max of the w1 and w3 scales then dequant and requant each expert. + if layer.w13_weight_scale.dim() == 2: # Shape: (num_experts, 2) + from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant + + # Get the maximum scale across w1 and w3 for each expert + max_w13_scales = layer.w13_weight_scale.max(dim=1).values + + # Requantize each expert's weights using the combined scale + # w13_weight has shape (num_experts, 2 * intermediate_size_per_partition, hidden_size) + # where the first intermediate_size_per_partition rows are w1, the next are w3 + intermediate_size_per_partition = layer.w13_weight.shape[1] // 2 + for expert_id in range(layer.w13_weight.shape[0]): + start = 0 + for shard_id in range(2): # w1 and w3 + # Dequantize using the original scale for this shard + dq_weight = per_tensor_dequantize( + layer.w13_weight[expert_id][ + start : start + intermediate_size_per_partition, : + ], + layer.w13_weight_scale[expert_id][shard_id], + ) + # Requantize using the combined max scale + ( + layer.w13_weight[expert_id][ + start : start + intermediate_size_per_partition, : + ], + _, + ) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id]) + + start += intermediate_size_per_partition + + # Update the scale parameter to be per-expert instead of per-shard + layer.w13_weight_scale = Parameter(max_w13_scales, requires_grad=False) + else: + layer.w13_weight_scale = Parameter( + layer.w13_weight_scale.data, requires_grad=False + ) + + if hasattr(layer, "w2_weight_scale") and layer.w2_weight_scale is not None: + layer.w2_weight_scale = Parameter( + layer.w2_weight_scale.data, requires_grad=False + ) + if hasattr(layer, "w13_input_scale") and layer.w13_input_scale is not None: + layer.w13_input_scale = Parameter( + layer.w13_input_scale.max(), requires_grad=False + ) + if hasattr(layer, "w2_input_scale") and layer.w2_input_scale is not None: + layer.w2_input_scale = Parameter( + layer.w2_input_scale.max(), requires_grad=False + ) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + use_fp8_w8a8=True, + per_channel_quant=False, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + ) + + return self.runner.run(dispatch_output, quant_info) + + +class ModelOptFp4Config(QuantizationConfig): + """Config class for FP4.""" + + def __init__( + self, + is_checkpoint_nvfp4_serialized: bool = False, + kv_cache_quant_algo: str = None, + group_size: int = None, + exclude_modules: List[str] = None, + ) -> None: + self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized + if is_checkpoint_nvfp4_serialized: + logger.warning( + "Detected nvfp4 checkpoint. Please note that the " + "format is experimental and subject to change." + ) + self.group_size = group_size + self.kv_cache_quant_algo = kv_cache_quant_algo + self.exclude_modules = exclude_modules + + @classmethod + def get_name(cls) -> str: + return "modelopt_fp4" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half, torch.float8_e4m3fn] + + @classmethod + def get_min_capability(cls) -> int: + return 100 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["hf_quant_config.json"] + + @staticmethod + def common_group_size(cfg: dict) -> int: + """Return the unique group_size across the config; raise if missing/mismatched.""" + sizes = set() + + # Top-level and 'quantization' block + v = cfg.get("group_size") + if isinstance(v, int): + sizes.add(v) + q = cfg.get("quantization") + if isinstance(q, dict): + v = q.get("group_size") + if isinstance(v, int): + sizes.add(v) + + # config_groups: accept group-level or nested dicts (e.g., weights/input_activations) + for g in (cfg.get("config_groups") or {}).values(): + if isinstance(g, dict): + v = g.get("group_size") + if isinstance(v, int): + sizes.add(v) + for sub in g.values(): + if isinstance(sub, dict): + v = sub.get("group_size") + if isinstance(v, int): + sizes.add(v) + + if not sizes: + raise ValueError("No group_size found in config.") + if len(sizes) > 1: + raise ValueError(f"Inconsistent group_size values: {sorted(sizes)}") + return next(iter(sizes)) + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: + # Handle two different config formats: + # 1. hf_quant_config.json format: {"quantization": {"quant_algo": "NVFP4", ...}} + # 2. config.json quantization_config format: {"quant_algo": "NVFP4", ...} + # In future modelopt will deprecate hf_quant_config.json, and only keep config.json. + # For legacy reasons, we keep hf_quant_config.json for now. + + # Initialize variables + kv_cache_quant_algo = None + group_size = None + exclude_modules = [] + + # Try flat format first (config.json quantization_config - preferred format) + quant_method = config.get("quant_algo") + if quant_method is not None: + # Flat format (config.json quantization_config) + # Note: FP4 models in config.json format may not have all the detailed fields + # that are present in hf_quant_config.json, so we need to handle defaults + kv_cache_quant_algo = config.get("kv_cache_quant_algo") + if not kv_cache_quant_algo: + # For config.json format, derive from kv_cache_scheme if available + kv_cache_scheme = config.get("kv_cache_scheme") + if ( + kv_cache_scheme + and kv_cache_scheme.get("type") == "float" + and kv_cache_scheme.get("num_bits") == 8 + ): + kv_cache_quant_algo = "FP8" + else: + kv_cache_quant_algo = "auto" + + group_size = ModelOptFp4Config.common_group_size(config) + exclude_modules = config.get("ignore", []) + else: + # Fall back to nested format (hf_quant_config.json - legacy format) + try: + quant_config = cls.get_from_keys(config, ["quantization"]) + quant_method = quant_config["quant_algo"] + kv_cache_quant_algo = quant_config.get("kv_cache_quant_algo") + if not kv_cache_quant_algo: + kv_cache_quant_algo = "auto" + group_size = ModelOptFp4Config.common_group_size(config) + exclude_modules = quant_config.get("exclude_modules", []) + except (ValueError, KeyError): + raise ValueError( + "Cannot find 'quant_algo' in the model's quantization config. " + "Expected either flat format (config.json) or nested format (hf_quant_config.json)." + ) + + if not quant_method in ["FP8", "NVFP4"]: + raise ValueError( + f"ModelOpt currently only supports: FP8, NVFP4" + " quantizations in sglang. Please check the " + "quantization config for your model's configuration." + ) + is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method + + if not (group_size and kv_cache_quant_algo) or exclude_modules is None: + logger.warning( + f"group_size: {group_size}," + f"kv_cache_quant_algo: {kv_cache_quant_algo}," + f"exclude_modules: {exclude_modules}" + ) + raise ValueError( + "NVFP4 quantization requires group size and " + "kv_cache_quant_algo specified in the quantization config" + ) + return cls( + is_checkpoint_nvfp4_serialized, + kv_cache_quant_algo, + group_size, + exclude_modules, + ) + + def is_layer_excluded(self, prefix: str, exclude_modules: list): + import regex as re + + fused_patterns = ["q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"] + prefix_split = prefix.split(".") + for pattern in exclude_modules: + regex_str = pattern.replace(".", r"\.").replace("*", r".*") + pattern_split = pattern.split(".") + if re.fullmatch(regex_str, prefix): + return True + elif ( + pattern_split[-1] in fused_patterns + and pattern_split[-1] in prefix_split[-1] + ): + # Check if the last part of the excluded pattern is contained in the last part of the prefix + # This handles fused modules like fused_qkv_a_proj_with_mqa that contain q_a_proj and kv_a_proj_with_mqa + # e.g., model.layers.{i}.self_attn.{fused_weight_name} + assert len(prefix_split) == 5 and len(pattern_split) == 5 + return True + return False + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.moe.fused_moe_triton.layer import FlashInferFP4MoE + + if isinstance(layer, LinearBase): + if is_layer_skipped(prefix, self.exclude_modules) or self.is_layer_excluded( + prefix, self.exclude_modules + ): + return UnquantizedLinearMethod() + return ModelOptFp4LinearMethod(self) + if self.kv_cache_quant_algo and isinstance(layer, RadixAttention): + return ModelOptFp8KVCacheMethod(self) + elif isinstance(layer, FlashInferFP4MoE): + # FlashInferFP4MoE needs the same quantization method but with compatible attribute handling + return ModelOptNvFp4FusedMoEMethod(self) + elif isinstance(layer, FusedMoE): + return ModelOptNvFp4FusedMoEMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class ModelOptFp4LinearMethod(LinearMethodBase): + """Linear method for NVFP4. + Supports loading NVFP4 checkpoints with the following structure: + + |Tensor Name | datatype | shape | + |----------------------------------------------------| + |input_scale | torch.float32 | scalar | + |weight | NVFP4(SE2M1) | [1, X, y/2] | + |weight_scale | FP8-E4M3 | [X, Y] | + |weight_scale_2 | torch.float32 | scalar | + + The weights are quantized per block of 16 elements. + Args: quant_config: The ModelOpt quantization config. + """ + + def __init__(self, quant_config: ModelOptFp4Config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + del input_size, output_size + if not self.quant_config.is_checkpoint_nvfp4_serialized: + raise ValueError( + "NVFP4 quantization was selected, " + " dynamic quantization is not supported." + ) + + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + if input_size_per_partition % 16 != 0: + raise ValueError( + "Unsupported model when in features size is " "not multiple of 16" + ) + + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_nvfp4_serialized + else params_dtype + ) + + weight = ModelWeightParameter( + data=torch.empty( + # 2 fp4 data is packed in one uint8 in the input dimension + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + input_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + + layer.register_parameter("input_scale", input_scale) + + weight_scale_2 = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale_2", weight_scale_2) + + weight_scale = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // self.quant_config.group_size, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + input_scale_2 = layer.input_scale.max().to(torch.float32) + weight_scale_2 = layer.weight_scale_2.max().to(torch.float32) + layer.input_scale = Parameter(input_scale_2, requires_grad=False) + layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False) + layer.alpha = Parameter( + layer.input_scale * layer.weight_scale_2, requires_grad=False + ) + layer.input_scale_inv = Parameter( + (1 / input_scale_2).to(torch.float32), requires_grad=False + ) + + # Pad and blockwise interleave weight_scale + scales = layer.weight_scale + scale_ndim = scales.ndim + if scale_ndim == 2: + scales = scales.unsqueeze(0) + assert scales.ndim == 3 + B, M, K = scales.shape + round_up_multiple = lambda x, m: (x + m - 1) // m * m + M_padded = round_up_multiple(M, 128) + K_padded = round_up_multiple(K, 4) + padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype) + padded_scales[:B, :M, :K] = scales + batches, rows, cols = padded_scales.shape + assert rows % 128 == 0 + assert cols % 4 == 0 + padded_scales = padded_scales.reshape(batches, rows // 128, 4, 32, cols // 4, 4) + padded_scales = padded_scales.permute((0, 1, 4, 3, 2, 5)) + padded_scales = padded_scales.contiguous().cuda() + padded_scales = ( + padded_scales.reshape(M_padded, K_padded) + if scale_ndim == 2 + else padded_scales.reshape(B, M_padded, K_padded) + ) + layer.weight_scale_interleaved = Parameter(padded_scales, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + output_dtype = x.dtype + x_m, _ = x.shape + w_n, _ = layer.weight.shape + output_shape = [x_m, w_n] + + # Quantize BF16 or FP16 to (FP4 and interleaved block scale) + x_fp4, x_scale_interleaved = scaled_fp4_quant(x, layer.input_scale_inv) + + assert x_fp4.dtype == torch.uint8 + assert x_scale_interleaved.dtype == torch.float8_e4m3fn + assert layer.weight.dtype == torch.uint8 + assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn + assert layer.alpha.dtype == torch.float32 + + w = layer.weight + w_scale_interleaved = layer.weight_scale_interleaved + if enable_flashinfer_fp4_gemm: + w = layer.weight.T + w_scale_interleaved = layer.weight_scale_interleaved.T + out = fp4_gemm( + x_fp4, + w, + x_scale_interleaved, + w_scale_interleaved, + layer.alpha, + output_dtype, + ) + if bias is not None: + out = out + bias + return out.view(*output_shape) + + +class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): + """ + MoE Method for FP4 Quantization with Blockscales and PerTensorScales + Args: + quant_config: NVFP4 Quant Config + """ + + def __init__(self, quant_config: ModelOptFp4Config): + self.quant_config = quant_config + if not is_sm100_supported(): + raise ValueError( + "Current platform does not support NVFP4" + " quantization. Please use Blackwell and" + " above." + ) + self.enable_flashinfer_trtllm_moe = should_use_flashinfer_trtllm_moe() + self._cache_permute_indices = {} + + @property + def enable_flashinfer_cutlass_moe(self) -> bool: + from sglang.srt.layers.moe import get_moe_runner_backend + + """Access the global enable_flashinfer_cutlass_moe setting.""" + return get_moe_runner_backend().is_flashinfer_cutlass() + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + if not self.quant_config.is_checkpoint_nvfp4_serialized: + raise ValueError( + "NVFP4 quantization was selected, " + " dynamic quantization is not supported." + ) + + # TODO(ch-wan): check if this is needed + layer.intermediate_size_per_partition = intermediate_size_per_partition + layer.params_dtype = params_dtype + layer.quant_config = self.quant_config + + weight_dtype = torch.uint8 + weight_scale_dtype = torch.float8_e4m3fn + weight_loader = extra_weight_attrs.get("weight_loader") + # GEMM 1 + w13_weight = ModelWeightParameter( + data=torch.empty( + layer.num_local_experts, + 2 * intermediate_size_per_partition, + # 2 fp4 items are packed in the input dimension + hidden_size // 2, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=2, + weight_loader=weight_loader, + ) + layer.register_parameter("w13_weight", w13_weight) + + # GEMM 2 + w2_weight = ModelWeightParameter( + data=torch.empty( + layer.num_local_experts, + hidden_size, + # 2 fp4 items are packed in the input dimension + intermediate_size_per_partition // 2, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=2, + weight_loader=weight_loader, + ) + layer.register_parameter("w2_weight", w2_weight) + + w13_weight_scale = ModelWeightParameter( + data=torch.empty( + layer.num_local_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.quant_config.group_size, + dtype=weight_scale_dtype, + ), + input_dim=1, + output_dim=2, + weight_loader=weight_loader, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + + # Only use `swizzle_blockscale` for shapes, not for real content + layer.w13_blockscale_swizzled = Parameter( + self.swizzle_blockscale(layer.w13_weight_scale), requires_grad=False + ) + + w2_weight_scale = ModelWeightParameter( + data=torch.empty( + layer.num_local_experts, + hidden_size, + intermediate_size_per_partition // self.quant_config.group_size, + dtype=weight_scale_dtype, + ), + input_dim=1, + output_dim=2, + weight_loader=weight_loader, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + layer.w2_blockscale_swizzled = Parameter( + self.swizzle_blockscale(layer.w2_weight_scale), requires_grad=False + ) + + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + ) + + w13_weight_scale_2 = PerTensorScaleParameter( + data=torch.empty(layer.num_local_experts, 2, dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2) + + w2_weight_scale_2 = PerTensorScaleParameter( + data=torch.empty(layer.num_local_experts, dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2) + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + + w13_input_scale = PerTensorScaleParameter( + data=torch.empty(layer.num_local_experts, 2, dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + + w2_input_scale = PerTensorScaleParameter( + data=torch.empty(layer.num_local_experts, dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + + def swizzle_blockscale(self, scale: torch.Tensor): + assert scale.dtype == torch.float8_e4m3fn + # Pad and blockwise interleave weight_scale + scale_ndim = scale.ndim + if scale.ndim == 2: + scale = scale.unsqueeze(0) + assert scale.ndim == 3 + B, M, K = scale.shape + round_up_multiple = lambda x, m: (x + m - 1) // m * m + M_padded = round_up_multiple(M, 128) + K_padded = round_up_multiple(K, 4) + padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype) + padded_scale[:B, :M, :K] = scale + batches, rows, cols = padded_scale.shape + assert rows % 128 == 0 + assert cols % 4 == 0 + padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32, cols // 4, 4) + swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5)) + swizzled_scale = swizzled_scale.contiguous().cuda() + return ( + swizzled_scale.reshape(M_padded, K_padded) + if scale_ndim == 2 + else swizzled_scale.reshape(B, M_padded, K_padded) + ) + + def prepare_static_weights_for_kernel( + self, + # args_dequant, + # args, + gemm1_weights, + gemm2_weights, + gemm1_scales_linear_fp4_bytes, + gemm2_scales_linear_fp4_bytes, + hidden_size, + intermediate_size, + num_experts, + ): + from flashinfer import ( + RoutingMethodType, + e2m1_and_ufp8sf_scale_to_float, + fp4_quantize, + next_positive_power_of_2, + nvfp4_block_scale_interleave, + reorder_rows_for_gated_act_gemm, + shuffle_matrix_a, + shuffle_matrix_sf_a, + ) + from flashinfer.fused_moe.core import ( + _maybe_get_cached_w2_permute_indices, + _maybe_get_cached_w3_w1_permute_indices, + ) + + """Prepare quantized weights for kernel (done offline with weights).""" + epilogue_tile_m = 128 # FIXME: this depends on the kernel internals + + # Convert quantized weights to proper formats + gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape( + num_experts, 2 * intermediate_size, hidden_size // 2 + ) # packed fp4 + gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view( + torch.float8_e4m3fn + ).reshape( + num_experts, 2 * intermediate_size, hidden_size // 16 + ) # fp8 scaling factors + + gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape( + num_experts, hidden_size, intermediate_size // 2 + ) # packed fp4 + gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view( + torch.float8_e4m3fn + ).reshape( + num_experts, hidden_size, intermediate_size // 16 + ) # fp8 scaling factors + + gemm1_weights_fp4_shuffled = [] + gemm1_scales_fp4_shuffled = [] + gemm2_weights_fp4_shuffled = [] + gemm2_scales_fp4_shuffled = [] + for i in range(num_experts): + # Calculate the permute indices for the following: + # 1. Reorder rows of W1 and scales for fused gated activation + # 2. Shuffle weights and scaling factors for transposed mma output + # for both w3_w1 and w2 weights and scale factors + permute_indices = _maybe_get_cached_w3_w1_permute_indices( + self._cache_permute_indices, + gemm1_weights_fp4[i].view(torch.uint8), + epilogue_tile_m, + ) + gemm1_weights_fp4_shuffled.append( + gemm1_weights_fp4[i] + .view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)] + .contiguous() + ) + + permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices( + self._cache_permute_indices, + gemm1_scales_linear_fp4[i].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + ) + gemm1_scales_fp4_shuffled.append( + nvfp4_block_scale_interleave( + gemm1_scales_linear_fp4[i] + .view(torch.uint8)[ + permute_sf_indices.to(gemm1_scales_linear_fp4.device) + ] + .contiguous() + ) + ) + + permute_indices = _maybe_get_cached_w2_permute_indices( + self._cache_permute_indices, + gemm2_weights_fp4[i].view(torch.uint8), + epilogue_tile_m, + ) + gemm2_weights_fp4_shuffled.append( + gemm2_weights_fp4[i] + .view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)] + .contiguous() + ) + + permute_sf_indices = _maybe_get_cached_w2_permute_indices( + self._cache_permute_indices, + gemm2_scales_linear_fp4[i].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + ) + gemm2_scales_fp4_shuffled.append( + nvfp4_block_scale_interleave( + gemm2_scales_linear_fp4[i] + .view(torch.uint8)[ + permute_sf_indices.to(gemm2_scales_linear_fp4.device) + ] + .contiguous() + ) + ) + + # Stack weights for all experts + gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled) + gemm1_scales_fp4_shuffled = ( + torch.stack(gemm1_scales_fp4_shuffled) + .view(torch.float8_e4m3fn) + .reshape(num_experts, 2 * intermediate_size, hidden_size // 16) + ) + + gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled) + gemm2_scales_fp4_shuffled = ( + torch.stack(gemm2_scales_fp4_shuffled) + .view(torch.float8_e4m3fn) + .reshape(num_experts, hidden_size, intermediate_size // 16) + ) + return ( + gemm1_weights_fp4_shuffled, + gemm1_scales_fp4_shuffled, + gemm2_weights_fp4_shuffled, + gemm2_scales_fp4_shuffled, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Process FP4 MoE weights after loading from serialized checkpoint. + + Only supports pre-quantized checkpoints with FP8 weights and scales. + """ + + # GEMM 1 scale processing + if not torch.allclose( + layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1] + ): + logger.warning_once( + "w1_weight_scale_2 must match w3_weight_scale_2. " + "Accuracy may be affected." + ) + + w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0] + layer.w13_weight_scale_2 = Parameter(w13_weight_scale_2, requires_grad=False) + + # Calculate input scales based on strategy + if self.enable_flashinfer_cutlass_moe or self.enable_flashinfer_trtllm_moe: + w13_input_scale = layer.w13_input_scale.max().to(torch.float32) + w2_input_scale = layer.w2_input_scale.max().to(torch.float32) + else: + w13_input_scale = layer.w13_input_scale.max(dim=1).values.to(torch.float32) + w2_input_scale = layer.w2_input_scale + + # Create shared parameters + layer.g1_alphas = Parameter( + (w13_input_scale * w13_weight_scale_2).to(torch.float32), + requires_grad=False, + ) + layer.g2_alphas = Parameter( + (w2_input_scale * layer.w2_weight_scale_2).to(torch.float32), + requires_grad=False, + ) + layer.w13_input_scale_quant = Parameter( + (1 / w13_input_scale).to(torch.float32), requires_grad=False + ) + layer.w2_input_scale_quant = Parameter( + (1 / w2_input_scale).to(torch.float32), requires_grad=False + ) + + # Validate weight scales + for name, weight_scale in [ + ("w13", layer.w13_weight_scale), + ("w2", layer.w2_weight_scale), + ]: + assert ( + weight_scale.shape[2] % 16 == 0 + ), f"Expected {name}_weight_scale.dim(2) to be divisible by 16" + assert ( + weight_scale.dtype == torch.float8_e4m3fn + ), f"{name} Weight Blockscale must be represented as FP8-E4M3" + + # Weight processing based on strategy + if ( + self.enable_flashinfer_trtllm_moe + and reorder_rows_for_gated_act_gemm is not None + and shuffle_matrix_sf_a is not None + ): + # FlashInfer TRTLLM processing - handles both w13 and w2 + ( + gemm1_weights_fp4_shuffled, + gemm1_scales_fp4_shuffled, + gemm2_weights_fp4_shuffled, + gemm2_scales_fp4_shuffled, + ) = self.prepare_static_weights_for_kernel( + layer.w13_weight, + layer.w2_weight, + layer.w13_weight_scale, + layer.w2_weight_scale, + layer.w2_weight.size(-2), # hidden_size + layer.w13_weight.size(-2) // 2, # intermediate_size + layer.w13_weight.size(0), # num_experts + ) + + # Set flashinfer parameters + layer.gemm1_weights_fp4_shuffled = Parameter( + gemm1_weights_fp4_shuffled, requires_grad=False + ) + layer.gemm2_weights_fp4_shuffled = Parameter( + gemm2_weights_fp4_shuffled, requires_grad=False + ) + layer.gemm1_scales_fp4_shuffled = Parameter( + gemm1_scales_fp4_shuffled, requires_grad=False + ) + layer.gemm2_scales_fp4_shuffled = Parameter( + gemm2_scales_fp4_shuffled, requires_grad=False + ) + + # Additional parameter needed for TRT-LLM + layer.g1_scale_c = Parameter( + (layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32), + requires_grad=False, + ) + + # Clean up weights that won't be used by TRT-LLM + del ( + layer.w2_weight, + layer.w2_weight_scale, + layer.w13_weight, + layer.w13_weight_scale, + ) + + else: + # CUTLASS processing - handle w13 and w2 separately + + # Process w13 weights + w13_blockscale_swizzled = self.swizzle_blockscale(layer.w13_weight_scale) + del layer.w13_weight_scale + layer.w13_blockscale_swizzled.data.copy_(w13_blockscale_swizzled) + layer.w13_weight = Parameter(layer.w13_weight.data, requires_grad=False) + + # Process w2 weights + w2_blockscale_swizzled = self.swizzle_blockscale(layer.w2_weight_scale) + del layer.w2_weight_scale + layer.w2_blockscale_swizzled.data.copy_(w2_blockscale_swizzled) + layer.w2_weight = Parameter(layer.w2_weight.data, requires_grad=False) + + # Both flashinfer cutlass and regular cutlass use same processing for w2 + + # Set up CUTLASS MoE parameters + device = layer.w13_weight.device + layer.cutlass_moe_params = CutlassMoEParams( + CutlassMoEType.BlockscaledFP4, + device, + num_experts=layer.num_experts, # global num experts + intermediate_size_per_partition=layer.w2_weight.shape[2] * 2, # n + hidden_size=layer.w13_weight.shape[2] * 2, + ) # k + + @property + def load_up_proj_weight_first(self) -> bool: + # FlashInfer CUTLASS kernel assumes [Up, Gate] Proj as W13 + return self.enable_flashinfer_cutlass_moe + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer: FusedMoE, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + assert ( + self.moe_runner_config.activation == "silu" + ), "Only SiLU activation is supported." + + moe_runner_config = self.moe_runner_config + + # Check if this is a FlashInferFP4MoE layer that should handle its own forward + if hasattr(layer, "gemm1_weights_fp4_shuffled"): + # This layer was processed with flashinfer TRTLLM - delegate to its own forward + return StandardCombineInput(hidden_states=layer.forward(x, topk_output)) + + if self.enable_flashinfer_cutlass_moe: + assert ( + not moe_runner_config.apply_router_weight_on_input + ), "apply_router_weight_on_input is not supported for Flashinfer" + # TRTLLM Cutlass moe takes in activations in BF16/Half/nvfp4 precision + # and fp4 quantized weights loaded from the checkpoint + topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids + + output_dtype = x.dtype + x_sf = None + if should_use_flashinfer_cutlass_moe_fp4_allgather(): + from flashinfer import fp4_quantize, nvfp4_block_scale_interleave + + # Quantize before comm, swizzle after. + if x.shape[0] > 0: + x, x_sf = fp4_quantize( + x, layer.w13_input_scale_quant, is_sf_swizzled_layout=False + ) + else: + x_col = x.shape[1] + x = torch.zeros(0, x_col // 2, dtype=torch.uint8, device=x.device) + x_sf = torch.zeros( + 0, x_col // 16, dtype=torch.uint8, device=x.device + ) + topk_weights, topk_ids, x, x_sf = get_tp_group().all_gatherv( + [topk_weights, topk_ids, x, x_sf], sizes=get_dp_global_num_tokens() + ) + x_sf = nvfp4_block_scale_interleave(x_sf) + + output = flashinfer_cutlass_fused_moe( + input=x, + token_selected_experts=topk_ids.to(torch.int), + token_final_scales=topk_weights, + fc1_expert_weights=layer.w13_weight.view(torch.long), + fc2_expert_weights=layer.w2_weight.view(torch.long), + output_dtype=output_dtype, + input_sf=x_sf, + quant_scales=[ + layer.w13_input_scale_quant, + layer.w13_blockscale_swizzled.view(torch.int32), + layer.g1_alphas, + layer.w2_input_scale_quant, + layer.w2_blockscale_swizzled.view(torch.int32), + layer.g2_alphas, + ], + ep_size=layer.moe_ep_size, + ep_rank=layer.moe_ep_rank, + tp_size=layer.moe_tp_size, + tp_rank=layer.moe_tp_rank, + tune_max_num_tokens=next_power_of_2(x.shape[0]), + )[0] + if should_use_flashinfer_cutlass_moe_fp4_allgather(): + output, global_output = get_local_dp_buffer(), output + get_tp_group().reduce_scatterv( + global_output, output=output, sizes=get_dp_global_num_tokens() + ) + return StandardCombineInput(hidden_states=output) + + from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4 + + topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids + output = cutlass_moe_fp4( + a=x, + a1_gscale=layer.w13_input_scale_quant, + w1_fp4=layer.w13_weight, + w1_blockscale=layer.w13_blockscale_swizzled, + w1_alphas=layer.g1_alphas, + a2_gscale=layer.w2_input_scale_quant, + w2_fp4=layer.w2_weight, + w2_blockscale=layer.w2_blockscale_swizzled, + w2_alphas=layer.g2_alphas, + topk_weights=topk_weights, + topk_ids=topk_ids, + params=layer.cutlass_moe_params, + apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, + ).to(x.dtype) + # Scale by routed_scaling_factor is fused into select_experts. + + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/moe_wna16.py b/python/sglang/srt/layers/quantization/moe_wna16.py new file mode 100644 index 000000000..531e4271f --- /dev/null +++ b/python/sglang/srt/layers/quantization/moe_wna16.py @@ -0,0 +1,501 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/moe_wna16.py +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import numpy as np +import torch + +from sglang.srt.distributed import get_tensor_model_parallel_rank +from sglang.srt.distributed.parallel_state import get_tp_group +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.quantization.awq import AWQConfig +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.gptq import GPTQConfig, GPTQMarlinConfig +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.utils import get_device_capability, set_weight_attrs + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + + +def get_weight_perm(num_bits: int): + perm_list: List[int] = [] + for i in range(32): + perm1: List[int] = [] + col = i // 4 + for block in [0, 1]: + for row in [ + 2 * (i % 4), + 2 * (i % 4) + 1, + 2 * (i % 4 + 4), + 2 * (i % 4 + 4) + 1, + ]: + perm1.append(16 * row + col + 8 * block) + for j in range(4): + perm_list.extend([p + 256 * j for p in perm1]) + + perm = np.array(perm_list) + + if num_bits == 4: + interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = np.array([0, 2, 1, 3]) + else: + raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) + + perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() + perm = torch.from_numpy(perm) + return perm + + +class MoeWNA16Config(QuantizationConfig): + """Config class for MOE WNA16 (W8A16/W4A16) quantization.""" + + def __init__( + self, + linear_quant_method: str, + weight_bits: int, + group_size: int, + has_zp: bool, + lm_head_quantized: bool, + modules_to_not_convert: Optional[List[str]], + full_config: Dict[str, Any], + ) -> None: + super().__init__() + self.weight_bits = weight_bits + self.group_size = group_size + self.has_zp = has_zp + self.bit8_pack_factor = 8 // self.weight_bits + self.lm_head_quantized = lm_head_quantized + self.linear_quant_method = linear_quant_method + self.full_config = full_config + self.use_marlin = False + # Avoid circular import + + if self.linear_quant_method == "gptq": + self.use_marlin = GPTQMarlinConfig.is_gptq_marlin_compatible(full_config) + elif self.linear_quant_method == "awq": + capability_tuple = get_device_capability() + device_capability = ( + -1 + if capability_tuple is None + else capability_tuple[0] * 10 + capability_tuple[1] + ) + awq_min_capability = AWQConfig.get_min_capability() + if device_capability < awq_min_capability: + raise ValueError( + "The quantization method moe_wna16 + awq is not supported " + "for the current GPU. " + f"Minimum capability: {awq_min_capability}. " + f"Current capability: {device_capability}." + ) + else: + raise ValueError("moe_wna16 only support gptq and awq.") + + if modules_to_not_convert is None: + self.modules_to_not_convert = [] + else: + self.modules_to_not_convert = modules_to_not_convert + + @classmethod + def get_name(cls) -> str: + return "moe_wna16" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 70 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["quantize_config.json"] + + def get_scaled_act_names(self) -> List[str]: + raise NotImplementedError + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> MoeWNA16Config: + quant_method = cls.get_from_keys(config, ["quant_method"]) + weight_bits = cls.get_from_keys(config, ["bits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + if quant_method == "gptq": + has_zp = not cls.get_from_keys(config, ["sym"]) + modules_to_not_convert = [] + elif quant_method == "awq": + has_zp = cls.get_from_keys(config, ["zero_point"]) + modules_to_not_convert = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + else: + raise ValueError("moe_wna16 only support gptq and awq.") + + return cls( + quant_method, + weight_bits, + group_size, + has_zp, + lm_head_quantized, + modules_to_not_convert, + config, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + if user_quant == "moe_wna16" and cls.is_moe_wna16_compatible(hf_quant_cfg): + return cls.get_name() + return None + + @classmethod + def is_moe_wna16_compatible(cls, quant_config: Dict[str, Any]): + # Extract data from quant config. + quant_method = quant_config.get("quant_method", "").lower() + num_bits = quant_config.get("bits") + desc_act = quant_config.get("desc_act") + + capability_tuple = get_device_capability() + device_capability = ( + -1 + if all(capability is None for capability in capability_tuple) + else capability_tuple[0] * 10 + capability_tuple[1] + ) + # Avoid circular import + awq_min_capability = AWQConfig.get_min_capability() + + gptq_compatible = quant_method == "gptq" and not desc_act and num_bits in [4, 8] + awq_compatible = ( + quant_method == "awq" + and num_bits == 4 + and device_capability >= awq_min_capability + ) + + return gptq_compatible or awq_compatible + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + # avoid circular import + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + + if is_layer_skipped_quant(prefix, self.modules_to_not_convert): + return UnquantizedLinearMethod() + elif isinstance(layer, LinearBase): + + if self.linear_quant_method == "gptq": + if self.use_marlin: + return GPTQMarlinConfig.from_config( + self.full_config + ).get_quant_method(layer, prefix) + else: + return GPTQConfig.from_config(self.full_config).get_quant_method( + layer, prefix + ) + elif self.linear_quant_method == "awq": + return AWQConfig.from_config(self.full_config).get_quant_method( + layer, prefix + ) + else: + raise ValueError("moe_wna16 only support gptq and awq.") + elif isinstance(layer, FusedMoE): + return MoeWNA16Method(self) + return None + + +def is_layer_skipped_quant(prefix: str, modules_to_not_convert: List[str]): + return any(module_name in prefix for module_name in modules_to_not_convert) + + +class MoeWNA16Method(FusedMoEMethodBase): + """Linear method for MOE WNA16 (W8A16/W4A16) quantization. + + Args: + quant_config: The MOE WNA16 (W8A16/W4A16) quantization config. + """ + + def __init__(self, quant_config: MoeWNA16Config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + layer.quant_config = self.quant_config + bit8_pack_factor = self.quant_config.bit8_pack_factor + group_size = self.quant_config.group_size + group_size_div_factor = 1 + + # make intermediate_size and hidden_size diviable by group_size + # we reduce the group size to ensure that + # and we would repeat the loaded_weight later + while intermediate_size_per_partition % group_size or hidden_size % group_size: + group_size = group_size // 2 + group_size_div_factor *= 2 + assert group_size >= 32 + layer.group_size = group_size + layer.group_size_div_factor = group_size_div_factor + + strategy = FusedMoeWeightScaleSupported.GROUP.value + extra_weight_attrs.update({"quant_method": strategy, "is_transposed": False}) + + assert "weight_loader" in extra_weight_attrs + weight_loader = extra_weight_attrs["weight_loader"] + wrapped_weight_loader = MoeWNA16Method.get_weight_loader(layer, weight_loader) + extra_weight_attrs["weight_loader"] = wrapped_weight_loader + + # Fused gate_up_proj (column parallel) + w13_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // bit8_pack_factor, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qweight", w13_qweight) + set_weight_attrs(w13_qweight, extra_weight_attrs) + + # down_proj (row parallel) + w2_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // bit8_pack_factor, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qweight", w2_qweight) + set_weight_attrs(w2_qweight, extra_weight_attrs) + + w13_scales = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // group_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_scales", w13_scales) + set_weight_attrs(w13_scales, extra_weight_attrs) + + w2_scales = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size, + intermediate_size_per_partition // group_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_scales", w2_scales) + set_weight_attrs(w2_scales, extra_weight_attrs) + + if self.quant_config.has_zp: + w13_qzeros = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition // bit8_pack_factor, + hidden_size // group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qzeros", w13_qzeros) + set_weight_attrs(w13_qzeros, extra_weight_attrs) + + w2_qzeros = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size // bit8_pack_factor, + intermediate_size_per_partition // group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qzeros", w2_qzeros) + set_weight_attrs(w2_qzeros, extra_weight_attrs) + + if self.quant_config.linear_quant_method == "gptq": + # some param are unused, but we need to init them in order to + # load weights + invalid_param_keys = ["w13_g_idx", "w2_g_idx"] + if not self.quant_config.has_zp: + invalid_param_keys += ["w13_qzeros", "w2_qzeros"] + for key in invalid_param_keys: + param = torch.nn.Parameter( + torch.empty((0,), dtype=torch.int32), requires_grad=False + ) + layer.register_parameter(key, param) + set_weight_attrs(param, extra_weight_attrs) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + assert ( + self.moe_runner_config.activation == "silu" + ), "Only SiLU activation is supported." + + weight_bits = self.quant_config.weight_bits + has_zp = self.quant_config.has_zp + + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_qweight, + w2_weight=layer.w2_qweight, + use_int4_w4a16=weight_bits == 4, + use_int8_w8a16=weight_bits == 8, + w13_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w13_zp=layer.w13_qzeros if has_zp else None, + w2_zp=layer.w2_qzeros if has_zp else None, + block_shape=[0, layer.group_size], + ) + return self.runner.run(dispatch_output, quant_info) + + @staticmethod + def get_weight_loader(layer, weight_loader): + + def convert_awq_tensor(tensor, tensor_type): + # convert awq qweight/qzeros to a standard format (assume int4) + # qweight: (k, n // pack_factor_bit32) -> (n, k // pack_factor_bit8) + # qzeros: (k // group_size, n // pack_factor_bit32) -> + # (n // pack_factor_bit8, k // group_size) + # pack_factor_bit32 = 32 // weight_bits + # pack_factor_bit8 = 8 // weight_bits + + # 0. suppose origin shape (a, b), dtype int32 + # 1. convert to uint8, shape (a, b) -> (a, 4 * b) + size0 = tensor.size(0) + tensor = tensor.view(torch.uint8) + + # 2. unpack to uint4 (only when weight_bits == 4) + # shape (a, 4 * b) -> (a, 4 * b, 2) + shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device) + tensor = (tensor[:, :, None] >> shifter) & 0xF + + # 3. change order, see + # https://github.com/casper-hansen/AutoAWQ/blob/v0.2.8/awq/utils/quant_utils.py + # shape -> (a, 4 * b * pack_factor_bit8) + reverse_awq_pack_order = [0, 4, 1, 5, 2, 6, 3, 7] + tensor = tensor.view(-1, 8)[:, reverse_awq_pack_order] + tensor = tensor.view(size0, -1) + + # 4. transpose, shape -> (4 * b * pack_factor_bit8, a) + tensor = tensor.T.contiguous() + + # 5. repack (only when weight_bits == 4) + # qweight shape -> (4 * b * pack_factor_bit8, a // pack_factor_bit8) + # qzeros shape -> (4 * b, a) + + if tensor_type == "qweight": + tensor = tensor[:, 1::2] * 16 + tensor[:, ::2] + elif tensor_type == "qzeros": + tensor = tensor[1::2, :] * 16 + tensor[::2, :] + return tensor + + def convert_gptq_int4_qzeros(tensor): + tensor = tensor.view(torch.uint8) + shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device) + tensor = (tensor[:, :, None] >> shifter) & 0xF + tensor = tensor + 1 + tensor = tensor[:, :, 0] + tensor[:, :, 1] * 16 + return tensor + + def moe_wna16_weight_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + ): + if "g_idx" in weight_name: + return + if not layer.quant_config.has_zp and "qzeros" in weight_name: + return + + device = get_tp_group().device + tp_rank = get_tensor_model_parallel_rank() + loaded_weight = loaded_weight.to(device) + shard_size = layer.intermediate_size_per_partition + + # convert gptq and awq weight to a standard format + if layer.quant_config.linear_quant_method == "awq": + assert layer.quant_config.weight_bits == 4 + if "weight" in weight_name: + loaded_weight = convert_awq_tensor(loaded_weight, "qweight") + elif "zeros" in weight_name: + loaded_weight = convert_awq_tensor(loaded_weight, "qzeros") + else: + loaded_weight = loaded_weight.T + elif layer.quant_config.linear_quant_method == "gptq": + assert layer.quant_config.weight_bits in [4, 8] + if "weight" in weight_name: + loaded_weight = loaded_weight.T.contiguous().view(torch.uint8) + elif "zeros" in weight_name: + # add 1 to gptq qzeros to align with awq + loaded_weight = loaded_weight.view(torch.uint8) + if layer.quant_config.weight_bits == 4: + loaded_weight = convert_gptq_int4_qzeros(loaded_weight).T + else: + loaded_weight = loaded_weight.T + 1 + else: + loaded_weight = loaded_weight.T + + # repeat the qzeros/scales to fit new group size + if ( + layer.group_size_div_factor > 1 + and "qzeros" in weight_name + or "scales" in weight_name + ): + loaded_weight = loaded_weight.repeat_interleave( + layer.group_size_div_factor, 1 + ) + + if "w13_qzeros" in weight_name: + tensor = loaded_weight.view( + layer.moe_tp_size, -1, loaded_weight.size(1) + )[tp_rank] + if shard_id == "w1": + param.data[expert_id, : shard_size // 2] = tensor + else: + param.data[expert_id, shard_size // 2 :] = tensor + elif "w2_qzeros" in weight_name: + param.data[expert_id] = loaded_weight.view( + loaded_weight.size(0), layer.moe_tp_size, -1 + )[:, tp_rank] + else: + weight_loader(param, loaded_weight, weight_name, shard_id, expert_id) + + return moe_wna16_weight_loader diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py new file mode 100644 index 000000000..0d98d00d6 --- /dev/null +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -0,0 +1,862 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/mxfp4.py + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, List, Optional + +import torch +from torch.nn.parameter import Parameter + +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.moe.utils import get_moe_runner_backend +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.utils import is_layer_skipped +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.utils import ( + direct_register_custom_op, + is_cuda, + is_flashinfer_available, + is_hip, + is_sm100_supported, + is_triton_kernels_available, + log_info_on_rank0, + mxfp_supported, + next_power_of_2, + round_up, + set_weight_attrs, +) + +_is_sm100_supported = is_cuda() and is_sm100_supported() +has_triton_kernels = is_triton_kernels_available() + + +if is_flashinfer_available(): + from flashinfer import ( + mxfp8_quantize, + shuffle_matrix_a, + shuffle_matrix_sf_a, + trtllm_fp4_block_scale_moe, + ) + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +_is_hip = is_hip() + +if _is_hip: + # import aiter + try: + from aiter import ActivationType, QuantType, dtypes + from aiter.fused_moe import fused_moe + from aiter.ops.triton.quant import dynamic_mxfp4_quant + from aiter.utility.fp4_utils import e8m0_shuffle + except ImportError as err: + ActivationType = QuantType = dtypes = fused_moe = dynamic_mxfp4_quant = ( + e8m0_shuffle + ) = err + + +def _swizzle_mxfp4(quant_tensor, scale, num_warps): + """weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel""" + import triton_kernels.matmul_ogs_details.opt_flags as opt_flags + from triton_kernels.numerics import InFlexData + from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor + from triton_kernels.tensor_details import layout + + value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout( + mx_axis=1 + ) + scale_layout, scale_layout_opts = layout.make_default_matmul_mxfp4_w_scale_layout( + mx_axis=1, num_warps=num_warps + ) + if _is_sm100_supported: + constraints = { + "is_persistent": True, + "epilogue_subtile": 1, + } + opt_flags.update_opt_flags_constraints(constraints) + # transpose the tensor so that the quantization axis is on dim1 + quant_tensor = quant_tensor.transpose(-2, -1) + scale = scale.transpose(-2, -1) + quant_tensor = convert_layout( + wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts + ) + scale = convert_layout(wrap_torch_tensor(scale), scale_layout, **scale_layout_opts) + return quant_tensor, InFlexData(), scale + + +def _dequant_mxfp4( + x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype +) -> torch.Tensor: + try: + from quark.torch.kernel import mx + except ImportError as err: + raise ImportError( + "The package `amd-quark` is required to use " + "MX-FP4 models. Please install it with `pip install " + "amd-quark`." + ) from err + + return mx.dq_mxfp4(x, scale, float_dtype) + + +def _dequant_mxfp4_fake( + x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype +) -> torch.Tensor: + return torch.empty( + (*x.shape[:-1], x.shape[-1] * 2), dtype=float_dtype, device=x.device + ) + + +def _quant_dequant_mxfp4( + x: torch.Tensor, scale_calculation_mode: str = "even" +) -> torch.Tensor: + try: + from quark.torch.kernel import mx + except ImportError as err: + raise ImportError( + "The package `amd-quark` is required to use " + "MX-FP4 models. Please install it with `pip install " + "amd-quark`." + ) from err + + return mx.qdq_mxfp4(x, scale_calculation_mode) + + +def _quant_dequant_mxfp4_fake( + x: torch.Tensor, scale_calculation_mode: str = "even" +) -> torch.Tensor: + return torch.empty_like(x) + + +direct_register_custom_op( + op_name="dequant_mxfp4", + op_func=_dequant_mxfp4, + mutates_args=[], + fake_impl=_dequant_mxfp4_fake, +) +dequant_mxfp4 = torch.ops.sglang.dequant_mxfp4 + +direct_register_custom_op( + op_name="quant_dequant_mxfp4", + op_func=_quant_dequant_mxfp4, + mutates_args=[], + fake_impl=_quant_dequant_mxfp4_fake, +) +quant_dequant_mxfp4 = torch.ops.sglang.quant_dequant_mxfp4 + + +class Mxfp4Config(QuantizationConfig): + + def __init__( + self, + ignored_layers: Optional[list[str]] = None, + is_checkpoint_mxfp4_serialized: bool = False, + ): + super().__init__() + self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized + self.ignored_layers = ignored_layers + + @classmethod + def from_config(cls, config): + + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_mxfp4_serialized = "mxfp4" in quant_method + + if _is_hip: + if mxfp_supported(): + return cls( + is_checkpoint_mxfp4_serialized=is_checkpoint_mxfp4_serialized + ) + else: + + platform = torch.cuda.get_device_properties(0).gcnArchName + raise ValueError( + f"Current platform {platform} not support mxfp4 computation" + ) + + return cls(is_checkpoint_mxfp4_serialized=is_checkpoint_mxfp4_serialized) + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_name(cls) -> str: + return "mxfp4" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def is_static_cfg(self): + return self.is_checkpoint_mxfp4_serialized + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["QuantizeMethodBase"]: + + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod + + if isinstance(layer, LinearBase): + if self.ignored_layers and is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + elif _is_hip: + return UnquantizedLinearMethod() + elif isinstance(layer, FusedMoE): + if self.is_checkpoint_mxfp4_serialized: + return Mxfp4MoEMethod(prefix=prefix) + else: + return Mxfp4DynamicQuantMoEMethod() + else: + if self.is_checkpoint_mxfp4_serialized: + raise NotImplementedError("Mxfp4 attention layer is not implemented") + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class Mxfp4MoEMethod(FusedMoEMethodBase): + + def __init__( + self, + prefix: str, + ): + super().__init__() + + self.prefix = prefix + self.topk_indices_dtype = None + self.use_triton_kernels = get_moe_runner_backend().is_triton_kernel() + self.with_bias = False + self.use_flashinfer = get_moe_runner_backend().is_flashinfer_mxfp4() + self.flashinfer_mxfp4_moe_precision = global_server_args_dict[ + "flashinfer_mxfp4_moe_precision" + ] + + self.triton_kernel_moe_forward = None + self.triton_kernel_moe_with_bias_forward = None + if torch.cuda.is_available() and has_triton_kernels: + from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import ( + triton_kernel_moe_forward as _tk_forward, + ) + from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import ( + triton_kernel_moe_with_bias_forward as _tk_with_bias_forward, + ) + + self.triton_kernel_moe_forward = _tk_forward + self.triton_kernel_moe_with_bias_forward = _tk_with_bias_forward + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + with_bias: bool = False, + **extra_weight_attrs, + ): + self.num_experts = num_experts + weight_dtype = torch.uint8 + scale_dtype = torch.uint8 + self.with_bias = with_bias + mxfp4_block = 32 + + # pad the intermediate size to be a multiple of 2 * mxfp4_block + # for to hold non-uniform sharded tensor as well as swizzling + intermediate_size_per_partition_after_pad = intermediate_size_per_partition + if _is_sm100_supported: + if self.use_flashinfer: + intermediate_size_per_partition_after_pad = round_up( + intermediate_size_per_partition, 256 + ) + hidden_size = round_up(hidden_size, 256) + else: + intermediate_size_per_partition_after_pad = round_up( + intermediate_size_per_partition, 64 + ) + elif has_triton_kernels: + # TODO: this is a hack to make + # intermediate_size_per_partition_after_pad the same as the + # per_rank_intermediate_size during weight loading + intermediate_size_per_partition_after_pad = round_up( + intermediate_size_per_partition, mxfp4_block + ) + + self.intermediate_size_per_partition = intermediate_size_per_partition_after_pad + + self.hidden_size = hidden_size + # Fused gate_up_proj (column parallel) + w13_weight = torch.nn.Parameter( + torch.zeros( + layer.num_local_experts, + 2 * intermediate_size_per_partition_after_pad, + hidden_size // 2, + dtype=weight_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w13_weight_scale = torch.nn.Parameter( + torch.zeros( + layer.num_local_experts, + 2 * intermediate_size_per_partition_after_pad, + hidden_size // mxfp4_block, + dtype=scale_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + w13_weight_bias = torch.nn.Parameter( + torch.zeros( + layer.num_local_experts, + 2 * intermediate_size_per_partition_after_pad, + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_bias", w13_weight_bias) + set_weight_attrs(w13_weight_bias, extra_weight_attrs) + + # down_proj (row parallel) + w2_weight = torch.nn.Parameter( + torch.zeros( + layer.num_local_experts, + hidden_size, + intermediate_size_per_partition_after_pad // 2, + dtype=weight_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w2_weight_scale = torch.nn.Parameter( + torch.zeros( + layer.num_local_experts, + hidden_size, + intermediate_size_per_partition_after_pad // mxfp4_block, + dtype=scale_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + w2_weight_bias = torch.nn.Parameter( + torch.zeros(layer.num_local_experts, hidden_size, dtype=torch.bfloat16), + requires_grad=False, + ) + layer.register_parameter("w2_weight_bias", w2_weight_bias) + set_weight_attrs(w2_weight_bias, extra_weight_attrs) + + def process_weights_after_loading(self, layer): + if self.use_flashinfer: + log_info_on_rank0( + logger, + f"Shuffling MoE weights for FlashInfer MXFP4 moe kernel (layer: {self.prefix}), it might take a while...", + ) + # TODO: these values are hardcoded for now, we need to get them from the model + layer.gemm1_alpha = Parameter( + torch.tensor([1.702] * self.num_experts, dtype=torch.float32).cuda(), + requires_grad=False, + ) + layer.gemm1_beta = Parameter( + torch.tensor([1.0] * self.num_experts, dtype=torch.float32).cuda(), + requires_grad=False, + ) + layer.gemm1_clamp_limit = Parameter( + torch.tensor([7.0] * self.num_experts, dtype=torch.float32).cuda(), + requires_grad=False, + ) + sf_block_size = 32 # mxfp4 block size + + assert ( + layer.w13_weight.dim() == 3 + and layer.w13_weight.shape[0] == self.num_experts + and layer.w13_weight.shape[1] + == self.intermediate_size_per_partition * 2 + and layer.w13_weight.shape[2] == self.hidden_size // 2 + ) + assert ( + layer.w13_weight_scale.dim() == 3 + and layer.w13_weight_scale.shape[0] == self.num_experts + and layer.w13_weight_scale.shape[1] + == self.intermediate_size_per_partition * 2 + and layer.w13_weight_scale.shape[2] == self.hidden_size // sf_block_size + ) + assert ( + layer.w2_weight.dim() == 3 + and layer.w2_weight.shape[0] == self.num_experts + and layer.w2_weight.shape[1] == self.hidden_size + and layer.w2_weight.shape[2] + == self.intermediate_size_per_partition // 2 + ) + assert ( + layer.w2_weight_scale.dim() == 3 + and layer.w2_weight_scale.shape[1] == self.hidden_size + and layer.w2_weight_scale.shape[2] + == self.intermediate_size_per_partition // sf_block_size + ) + assert ( + layer.w13_weight_bias.dim() == 2 + and layer.w13_weight_bias.shape[0] == self.num_experts + and layer.w13_weight_bias.shape[1] + == self.intermediate_size_per_partition * 2 + ) + assert ( + layer.w2_weight_bias.dim() == 2 + and layer.w2_weight_bias.shape[0] == self.num_experts + and layer.w2_weight_bias.shape[1] == self.hidden_size + ) + + w13_weight_scale = layer.w13_weight_scale.data + w2_weight_scale = layer.w2_weight_scale.data + w13_weight = layer.w13_weight.data + w2_weight = layer.w2_weight.data + w13_bias = layer.w13_weight_bias.data.to(torch.float32) + w2_bias = layer.w2_weight_bias.data.to(torch.float32) + + # Swap w1 and w3 as the definition of + # swiglu is different in the trtllm-gen + def swap_every_two_rows(x, axis=-1): + shape = x.shape + if axis < 0: + axis = len(shape) + axis + + # Create a new shape with pairs swapped along specified axis + new_shape = list(shape) + new_shape[axis] = shape[axis] // 2 + new_shape.insert(axis + 1, 2) + + # Reshape to expose pairs, swap them, and reshape back + x = x.reshape(*new_shape) + x = x.flip(axis + 1) + new_shape = list(shape) + return x.reshape(*new_shape) + + w13_weight_scale = swap_every_two_rows(w13_weight_scale, -2) + w13_weight = swap_every_two_rows(w13_weight, -2) + w13_bias = swap_every_two_rows(w13_bias, -1) + + # Shuffle weights and scaling factors for transposed mma output + gemm1_weights_mxfp4_shuffled = [] + gemm1_scales_mxfp4_shuffled = [] + gemm2_weights_mxfp4_shuffled = [] + gemm2_scales_mxfp4_shuffled = [] + gemm1_bias_shuffled = [] + gemm2_bias_shuffled = [] + epilogue_tile_m = 128 # FIXME: this depends on the kernel internals + for i in range(self.num_experts): + gemm1_weights_mxfp4_shuffled.append( + shuffle_matrix_a(w13_weight[i].view(torch.uint8), epilogue_tile_m) + ) + gemm1_scales_mxfp4_shuffled.append( + shuffle_matrix_sf_a( + w13_weight_scale[i].view(torch.uint8), epilogue_tile_m + ) + ) + gemm1_bias_shuffled.append( + shuffle_matrix_a( + w13_bias[i].clone().reshape(-1, 1), epilogue_tile_m + ) + ) + + gemm2_weights_mxfp4_shuffled.append( + shuffle_matrix_a(w2_weight[i].view(torch.uint8), epilogue_tile_m) + ) + gemm2_scales_mxfp4_shuffled.append( + shuffle_matrix_sf_a( + w2_weight_scale[i].view(torch.uint8), epilogue_tile_m + ) + ) + gemm2_bias_shuffled.append( + shuffle_matrix_a(w2_bias[i].clone().reshape(-1, 1), epilogue_tile_m) + ) + + w13_weight = torch.stack(gemm1_weights_mxfp4_shuffled) + w13_weight_scale = ( + torch.stack(gemm1_scales_mxfp4_shuffled) + .reshape( + self.num_experts, + 2 * self.intermediate_size_per_partition, + self.hidden_size // sf_block_size, + ) + .view(torch.float8_e4m3fn) + ) + + w2_weight = torch.stack(gemm2_weights_mxfp4_shuffled) + w2_weight_scale = ( + torch.stack(gemm2_scales_mxfp4_shuffled) + .reshape( + self.num_experts, + self.hidden_size, + self.intermediate_size_per_partition // sf_block_size, + ) + .view(torch.float8_e4m3fn) + ) + + layer.w13_weight = Parameter(w13_weight, requires_grad=False) + layer.w13_weight_scale = Parameter(w13_weight_scale, requires_grad=False) + layer.w2_weight = Parameter(w2_weight, requires_grad=False) + layer.w2_weight_scale = Parameter(w2_weight_scale, requires_grad=False) + layer.w13_weight_bias = Parameter( + torch.stack(gemm1_bias_shuffled).reshape(self.num_experts, -1), + requires_grad=False, + ) + layer.w2_weight_bias = Parameter( + torch.stack(gemm2_bias_shuffled).reshape(self.num_experts, -1), + requires_grad=False, + ) + return + + if self.use_triton_kernels: + + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig + + w13_weight_bias = layer.w13_weight_bias.to(torch.float32) + w2_weight_bias = layer.w2_weight_bias.to(torch.float32) + + layer.w13_weight_bias = Parameter(w13_weight_bias, requires_grad=False) + layer.w2_weight_bias = Parameter(w2_weight_bias, requires_grad=False) + + num_warps = 8 + + w13_weight, w13_flex, w13_scale = _swizzle_mxfp4( + layer.w13_weight, layer.w13_weight_scale, num_warps + ) + w2_weight, w2_flex, w2_scale = _swizzle_mxfp4( + layer.w2_weight, layer.w2_weight_scale, num_warps + ) + + self.w13_precision_config = PrecisionConfig( + weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex) + ) + self.w2_precision_config = PrecisionConfig( + weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex) + ) + + self.w13_weight_triton_tensor = w13_weight + self.w2_weight_triton_tensor = w2_weight + del layer.w13_weight + del layer.w2_weight + else: + from triton_kernels.numerics_details.mxfp import upcast_from_mxfp + + w13_weight = upcast_from_mxfp( + layer.w13_weight, layer.w13_weight_scale, dtype=torch.bfloat16, axis=-1 + ) + w2_weight = upcast_from_mxfp( + layer.w2_weight, layer.w2_weight_scale, dtype=torch.bfloat16, axis=-1 + ) + del layer.w13_weight + del layer.w2_weight + del layer.w13_weight_scale + del layer.w2_weight_scale + layer.w13_weight = Parameter(w13_weight.data, requires_grad=False) + layer.w2_weight = Parameter(w2_weight.data, requires_grad=False) + torch.cuda.empty_cache() + + def _get_tile_tokens_dim(self, x: torch.Tensor, top_k: int): + # Number of tokens in the input tensor. + num_tokens = x.shape[0] + # Factor to account for the imbalance of the experts. + # factor equals to the + # max_real_num_tokens_per_expert / perfect_num_tokens_per_expert + # - 1.0 means perfect expert distribution. + # - > 1.0 means some experts have more + # tokens than the perfect distribution. + # - < 1.0 does not make sense. + imbalance_factor = 1.3 + # Calculate the number of tokens per expert + # assuming perfect distribution. + num_tokens_per_expert = (num_tokens * top_k) // self.num_experts + # Apply the imbalance factor. + num_tokens_per_expert = int(num_tokens_per_expert * imbalance_factor) + # And pad the number to the next power of 2. + tile_tokens_dim = next_power_of_2(num_tokens_per_expert) + # Cap to 8-64 tokens per CTA tile + # as it's the range supported by the kernel. + tile_tokens_dim = min(max(tile_tokens_dim, 8), 64) + + return tile_tokens_dim + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + from sglang.srt.layers.moe.topk import TopKOutputChecker + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + moe_runner_config = self.moe_runner_config + + if self.use_flashinfer: + # When bf16 mode is enabled, we don't need to quantize the input, + # TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations, + # which can theoretically improve performance + if self.flashinfer_mxfp4_moe_precision == "bf16": + assert x.dtype == torch.bfloat16 + x_quant = x + x_scale = None + + # May be fused later if this code branch is frequently needed + origin_hidden_states_dim = x_quant.shape[-1] + if self.hidden_size != origin_hidden_states_dim: + x_quant = torch.nn.functional.pad( + x_quant, + (0, self.hidden_size - origin_hidden_states_dim), + mode="constant", + value=0.0, + ) + elif self.flashinfer_mxfp4_moe_precision == "default": + x_quant, x_scale = mxfp8_quantize(x, False, alignment=self.hidden_size) + x_scale = x_scale.view(torch.float8_e4m3fn).reshape(-1) + else: + raise NotImplementedError + + assert x_quant.shape[-1] == self.hidden_size + assert TopKOutputChecker.format_is_bypassed(topk_output) + + top_k = topk_output.topk_config.top_k + router_logits = topk_output.router_logits + + trtllm_gen_output = trtllm_fp4_block_scale_moe( + router_logits.to(torch.bfloat16), + None, # routing_bias + x_quant, + x_scale, + layer.w13_weight, # uint8 (e2m1 x 2) + layer.w13_weight_scale, # uint8 (e4m3 x 2) + layer.w13_weight_bias, # fp32 per expert per channel + layer.gemm1_alpha, # fp32 per expert + layer.gemm1_beta, # fp32 per expert + layer.gemm1_clamp_limit, # fp32 per expert + layer.w2_weight, # uint8 (e2m1 x 2) + layer.w2_weight_scale, # ue8m0 + layer.w2_weight_bias, # fp32 per expert per channel + None, # output1_scale_scalar + None, # output1_scale_gate_scalar + None, # output2_scale_scalar + layer.num_experts, + top_k, + None, # n_group # TODO: support n_group + None, # topk_group # TODO: support topk_group + self.intermediate_size_per_partition, # padded to multiple of 256 + layer.moe_ep_rank * layer.num_local_experts, # local_expert_offset + layer.num_local_experts, # local num experts + None, + self._get_tile_tokens_dim(x, top_k), + 1, # routing_method_type, renormalize + True, # do finalize + )[0] + return StandardCombineInput(hidden_states=trtllm_gen_output) + + if self.use_triton_kernels: + assert ( + layer.moe_ep_size == 1 + ), "Expert parallel is not supported when using triton kernels" + if self.with_bias: + output = self.triton_kernel_moe_with_bias_forward( + hidden_states=x, + w1=self.w13_weight_triton_tensor, + w1_pcg=self.w13_precision_config, + w2=self.w2_weight_triton_tensor, + w2_pcg=self.w2_precision_config, + b1=layer.w13_weight_bias, + b2=layer.w2_weight_bias, + topk_output=topk_output, + moe_runner_config=moe_runner_config, + ) + else: + output = self.triton_kernel_moe_forward( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_output=topk_output, + moe_runner_config=moe_runner_config, + ) + return StandardCombineInput(hidden_states=output) + else: + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_bias=layer.w13_weight_bias, + w2_weight_bias=layer.w2_weight_bias, + ) + return self.runner.run(dispatch_output, quant_info) + + +class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase): + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # Allocate 2 scales for w1 and w3 respectively. + # They will be combined to a single scale after weight loading. + w13_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + # Add the quantization method used (per tensor/grouped/channel) + # to ensure the weight scales are loaded in properly + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + + layer.w13_input_scale = None + layer.w2_input_scale = None + + def mxfp4_quantize(self, w): + w_shape = w.shape + w_need_reshape = True if w.dim() != 2 else False + + if w_need_reshape: + w_last_dim_size = w_shape[-1] + w = w.view(-1, w_last_dim_size) + + w, mx_scales = dynamic_mxfp4_quant(w) + + if w_need_reshape: + w_new_shape = w_shape[:-1] + (w.shape[-1],) + w = w.view(w_new_shape) + + mx_scales = e8m0_shuffle(mx_scales) + + return w, mx_scales + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w13, w13_mx_scales = self.mxfp4_quantize(layer.w13_weight.data) + w2, w2_mx_scales = self.mxfp4_quantize(layer.w2_weight.data) + + layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter(w13_mx_scales, requires_grad=False) + + layer.w2_weight = torch.nn.Parameter(w2, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter(w2_mx_scales, requires_grad=False) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + topk_weights, topk_ids, _ = topk_output + if _is_hip: + topk_weights = topk_weights.to( + torch.float32 + ) # aiter's moe_sorting requires topk_weights to be FP32 + output = fused_moe( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + quant_type=QuantType.per_1x32, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + activation=( + ActivationType.Silu + if self.moe_runner_config.activation == "silu" + else ActivationType.Gelu + ), + doweight_stage1=False, + ) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/mxfp4_tensor.py b/python/sglang/srt/layers/quantization/mxfp4_tensor.py new file mode 100644 index 000000000..76cb92c54 --- /dev/null +++ b/python/sglang/srt/layers/quantization/mxfp4_tensor.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional + +import torch + + +# https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/modelopt/torch/quantization/qtensor/mxfp4_tensor.py +class MXFP4QuantizeUtil: + E2M1_max = 6.0 + + E2M1_values = [0, 0.5, 1, 1.5, 2, 3, 4, 6] + E2M1_bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5]) + + @classmethod + def quantize(cls, input: torch.Tensor, block_size: Optional[int]) -> tuple: + """Converting a tensor to a quantized format based on MXFP4 quantization. Only E4M3 is supported. + Args: + input (torch.Tensor): The input tensor to be quantized. + block_sizes (dict | None): The block sizes for quantization. + """ + + def cast_fp4(x): + sign = torch.sign(x) + sign_bit = (2 - sign) // 2 + ord_ = torch.sum( + (x.abs().unsqueeze(-1) - cls.E2M1_bounds.to(x.device)) > 0, dim=-1 + ) + fp4_val = (sign_bit * 0b1000 + ord_).to(torch.uint8) + return fp4_val + + def fuse_uint4_to_uint8(x): + # If the last dimension is odd, pad with zeros + # If this behavior is not desired, please modify the code accordingly + left_side = x[..., 0::2] # Even indices (0, 2, 4...) + right_side = x[..., 1::2] # Odd indices (1, 3, 5...) + new_data = ( + right_side.clone() << 4 + ) # Put odd indices (higher addresses) in high bits + new_data[ + ..., : left_side.shape[-1] + ] += left_side # Put even indices in low bits + return new_data + + if block_size is None: + block_size = 32 + + original_shape = input.shape + original_dtype = input.dtype + input = input.view(-1, block_size) + # get scales + input_amax = input.abs().max(dim=-1, keepdim=True).values + descale = input_amax / cls.E2M1_max + min_value = torch.tensor(-127.0, device=descale.device) + e8m0_scale = torch.ceil(torch.maximum(torch.log2(descale), min_value)) + + input = (input / torch.exp2(e8m0_scale)).view(original_shape) + input_q = cast_fp4(input) + input_q = fuse_uint4_to_uint8(input_q) + e8m0_scale = (e8m0_scale + 127).to(torch.uint8) + return cls(original_shape, original_dtype, input_q), e8m0_scale + + @classmethod + def dequantize(cls, quantized_data, dtype: torch.dtype, scale, block_sizes): + """Dequantze MXFP4 packed tensor to a target dtype.""" + + def unfuse_uint8_to_uint4(x): + """Unfuse uint8 values back to uint4 values. + This is the inverse operation of fuse_uint4_to_uint8. + """ + # Extract the lower 4 bits (even indices) + left_side = x & 0x0F + + # Extract the upper 4 bits (odd indices) + right_side = (x >> 4) & 0x0F + + # Create a new tensor with alternating values + shape = list(x.shape) + shape[-1] = shape[-1] * 2 + result = torch.zeros(shape, dtype=torch.uint8, device=x.device) + + # Fill in the values - even indices get low bits, odd indices get high bits + result[..., 0::2] = left_side # Even indices from low bits + result[..., 1::2] = right_side # Odd indices from high bits + + return result + + e8m0_scale = scale + block_size = block_sizes[-1] + + # Unfuse the uint8 values back to uint4 + x_unfused = unfuse_uint8_to_uint4(quantized_data) + # Extract sign and magnitude + sign = 1 - 2 * ((x_unfused & 0b1000) >> 3).to( + torch.float32 + ) # Extract sign bit and convert to +1/-1 + magnitude = x_unfused & 0b0111 # Extract magnitude bits + magnitude = magnitude.to(torch.long) + + # Create a tensor with the E2M1 values + values = torch.tensor(cls.E2M1_values, device=quantized_data.device) + + # Use gather to index the values tensor properly + # We need to reshape magnitude to match the dimensions we want to gather along + original_shape = magnitude.shape + x_float = values[magnitude.reshape(-1)].reshape(original_shape) + + # Apply sign and scale + x_float = sign.float() * x_float + + # Reshape to apply block-wise scaling + x_float = x_float.reshape(-1, block_size) + + # Apply the E8M0 scale + scale_factor = torch.exp2(e8m0_scale.float() - 127) + scale_factor = scale_factor.reshape(-1, 1) # Reshape for proper broadcasting + + # Apply scaling and reshape back to original shape + x_float = x_float * scale_factor + + # Reshape back to the original shape + return x_float.reshape(original_shape).to(dtype) diff --git a/python/sglang/srt/layers/quantization/petit.py b/python/sglang/srt/layers/quantization/petit.py new file mode 100644 index 000000000..2c608507c --- /dev/null +++ b/python/sglang/srt/layers/quantization/petit.py @@ -0,0 +1,252 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py + + +import logging +from typing import Any, Callable, Dict, List, Optional + +import regex as re +import torch +from torch.nn.parameter import Parameter + +from sglang.srt.layers.linear import LinearBase, UnquantizedLinearMethod +from sglang.srt.layers.parameter import ModelWeightParameter, PerTensorScaleParameter +from sglang.srt.layers.quantization.base_config import ( + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.petit_utils import ( + apply_petit_nvfp4_linear, + prepare_nvfp4_layer_for_petit, + verify_petit_nvfp4_supported, +) +from sglang.srt.layers.quantization.utils import is_layer_skipped +from sglang.srt.utils import is_hip + +_is_hip = is_hip() + +# Initialize logger for the module +logger = logging.getLogger(__name__) + + +# Configuration class to support the NVFP4 quantized model generated by the ModelOpt quantization tool +class PetitNvFp4Config(QuantizationConfig): + """Config class for Petit FP4.""" + + def __init__( + self, + is_checkpoint_nvfp4_serialized: bool = False, + kv_cache_quant_algo: str = None, + group_size: int = None, + exclude_modules: List[str] = None, + ) -> None: + self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized + if is_checkpoint_nvfp4_serialized: + logger.warning( + "Detected nvfp4 checkpoint. Please note that the " + "format is experimental and subject to change." + ) + self.group_size = group_size + self.kv_cache_quant_algo = kv_cache_quant_algo + self.exclude_modules = exclude_modules + + @classmethod + def get_name(cls) -> str: + return "petit_nvfp4" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + # Petit supports the gfx90a and gfx942 GPUs + return 90 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["hf_quant_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "PetitNvFp4Config": + quant_config = cls.get_from_keys(config, ["quantization"]) + quant_method = quant_config["quant_algo"] + group_size = quant_config.get("group_size", None) + verify_petit_nvfp4_supported(quant_method, group_size) + + is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method + kv_cache_quant_algo = quant_config["kv_cache_quant_algo"] + if not kv_cache_quant_algo: + kv_cache_quant_algo = "auto" + exclude_modules = quant_config.get("exclude_modules", None) + if not (group_size and kv_cache_quant_algo and (exclude_modules is not None)): + logger.warning( + f"group_size: {group_size}," + f"kv_cache_quant_algo: {kv_cache_quant_algo}," + f"exclude_modules: {exclude_modules}" + ) + raise ValueError( + "NVFP4 quantization requires group size and " + "kv_cache_quant_algo specified in " + "hf_quant_config.json" + ) + return cls( + is_checkpoint_nvfp4_serialized, + kv_cache_quant_algo, + group_size, + exclude_modules, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + can_convert = cls.is_petit_nvfp4_compatible(hf_quant_cfg) + if can_convert: + return cls.get_name() + return None + + @classmethod + def is_petit_nvfp4_compatible(cls, quant_config: Dict[str, Any]) -> bool: + quant_method = quant_config.get("quant_method", "").lower() + return _is_hip and quant_method == "modelopt" + + def is_layer_excluded(self, prefix: str, exclude_modules: list): + for pattern in exclude_modules: + regex_str = pattern.replace(".", r"\.").replace("*", r".*") + if re.fullmatch(regex_str, prefix): + return True + return False + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["QuantizeMethodBase"]: + if isinstance(layer, LinearBase): + if is_layer_skipped(prefix, self.exclude_modules) or self.is_layer_excluded( + prefix, self.exclude_modules + ): + return UnquantizedLinearMethod() + return PetitNvFp4LinearMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class PetitNvFp4LinearMethod(LinearMethodBase): + """Linear method for NVFP4. + Supports loading NVFP4 checkpoints with the following structure: + + |Tensor Name | datatype | shape | + |----------------------------------------------------| + |input_scale | torch.float32 | scalar | + |weight | NVFP4(SE2M1) | [1, X, y/2] | + |weight_scale | FP8-E4M3 | [X, Y] | + |weight_scale_2 | torch.float32 | scalar | + + The weights are quantized per block of 16 elements. + Args: quant_config: The ModelOpt quantization config. + """ + + def __init__(self, quant_config: PetitNvFp4Config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + del input_size, output_size + if not self.quant_config.is_checkpoint_nvfp4_serialized: + raise ValueError( + "NVFP4 quantization was selected, " + " dynamic quantization is not supported." + ) + + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + if input_size_per_partition % 16 != 0: + raise ValueError( + "Unsupported model when in features size is " "not multiple of 16" + ) + + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_nvfp4_serialized + else params_dtype + ) + + weight = ModelWeightParameter( + data=torch.empty( + # 2 fp4 data is packed in one uint8 in the input dimension + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + input_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + + layer.register_parameter("input_scale", input_scale) + + weight_scale_2 = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale_2", weight_scale_2) + + weight_scale = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // self.quant_config.group_size, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + input_scale_2 = layer.input_scale.max().to(torch.float32) + weight_scale_2 = layer.weight_scale_2.max().to(torch.float32) + layer.input_scale = Parameter(input_scale_2, requires_grad=False) + layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False) + layer.alpha = Parameter( + layer.input_scale * layer.weight_scale_2, requires_grad=False + ) + + prepare_nvfp4_layer_for_petit(layer) + del layer.input_scale + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_petit_nvfp4_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + weight_scale_2=layer.weight_scale_2, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) diff --git a/python/sglang/srt/layers/quantization/petit_utils.py b/python/sglang/srt/layers/quantization/petit_utils.py new file mode 100644 index 000000000..529869f24 --- /dev/null +++ b/python/sglang/srt/layers/quantization/petit_utils.py @@ -0,0 +1,104 @@ +from typing import Optional + +import torch + +try: + from petit_kernel import mul_nvfp4_a16, process_nvfp4_scales, repack_nvfp4 +except ImportError: + + def _check_petit_nvfp4_supported( + quant_method: str, group_size: Optional[int] + ) -> tuple[bool, Optional[str]]: + return ( + False, + "Petit is not installed. Please install it with `pip install petit-kernel`.", + ) + + def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None: + raise ValueError( + "Petit is not installed. Please install it with `pip install petit-kernel`." + ) + + def apply_petit_nvfp4_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_scale_2: torch.Tensor, + size_n: int, + size_k: int, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + raise ValueError( + "Petit is not installed. Please install it with `pip install petit-kernel`." + ) + + +def _check_petit_nvfp4_supported( + quant_method: str, group_size: Optional[int] +) -> tuple[bool, Optional[str]]: + if quant_method != "NVFP4": + return ( + False, + "Petit currently only supports: NVFP4" + " quantizations in sglang. Please check the " + "`hf_quant_config.json` file for your model's " + "quant configuration.", + ) + if group_size is not None and group_size != 16: + return ( + False, + "Petit currently only supports: group_size=16" " quantizations.", + ) + return (True, None) + + +def verify_petit_nvfp4_supported(quant_method: str, group_size: Optional[int]) -> None: + supported, error_msg = _check_petit_nvfp4_supported(quant_method, group_size) + if not supported: + raise ValueError(error_msg) + + +def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None: + # Repack weights to petit format + part_size_n = layer.output_size_per_partition + part_size_k = layer.input_size_per_partition + qweight = layer.weight.view(torch.int32).contiguous() + petit_qweight = repack_nvfp4(qweight, size_n=part_size_n, size_k=part_size_k) + layer.weight = torch.nn.Parameter(petit_qweight, requires_grad=False) + + # Permute scales + weight_scale = process_nvfp4_scales( + scales=layer.weight_scale, size_k=part_size_k, size_n=part_size_n + ) + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + + return + + +def apply_petit_nvfp4_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_scale_2: torch.Tensor, + size_n: int, + size_k: int, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + reshaped_x = input.reshape(-1, input.shape[-1]) + out_shape = input.shape[:-1] + (size_n,) + + # TODO: Use auto-tuning to find the performant solution_id + output = mul_nvfp4_a16( + a=reshaped_x, + b=weight, + s=weight_scale, + global_scale=weight_scale_2, + size_m=reshaped_x.size(0), + size_n=size_n, + size_k=size_k, + solution_id=-1, + ) + if bias is not None: + output.add_(bias) # In-place add + + return output.reshape(out_shape) diff --git a/python/sglang/srt/layers/quantization/qoq.py b/python/sglang/srt/layers/quantization/qoq.py new file mode 100644 index 000000000..ec0fda482 --- /dev/null +++ b/python/sglang/srt/layers/quantization/qoq.py @@ -0,0 +1,245 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import torch +from torch.nn.parameter import Parameter + +from sglang.srt.layers.parameter import ( + ChannelQuantScaleParameter, + GroupQuantScaleParameter, + ModelWeightParameter, +) +from sglang.srt.layers.quantization.base_config import ( + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8 +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() +if _is_cuda: + from sgl_kernel import qserve_w4a8_per_chn_gemm, qserve_w4a8_per_group_gemm + + +QoQ_SUPPORTED_WEIGHT_BITS = [4] +QoQ_SUPPORTED_GROUP_SIZES = [-1, 128] + + +class QoQConfig(QuantizationConfig): + """Config class for QoQ Quantization. + + - Weight: static, per-channel/group, asymmetric + - Activation: dynamic, per-token, symmetric + + Reference: https://arxiv.org/abs/2405.04532 + https://github.com/mit-han-lab/omniserve + """ + + def __init__(self, weight_bits: int, group_size: int) -> None: + self.weight_bits = weight_bits + self.group_size = group_size + + # Verify + if self.weight_bits not in QoQ_SUPPORTED_WEIGHT_BITS: + raise ValueError( + f"QoQ does not support weight_bits = {self.weight_bits}. " + f"Only weight_bits = {QoQ_SUPPORTED_WEIGHT_BITS} " + "are supported." + ) + if self.group_size not in QoQ_SUPPORTED_GROUP_SIZES: + raise ValueError( + f"QoQ does not support group_size = {self.group_size}. " + f"Only group_sizes = {QoQ_SUPPORTED_GROUP_SIZES} " + "are supported." + ) + + # 4 bits packed into 8 bit datatype. + self.pack_factor = 8 // self.weight_bits + + def __repr__(self) -> str: + return "QoQConfig(weight_bits={}, group_size={})".format( + self.weight_bits, self.group_size + ) + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_name(cls) -> str: + return "qoq" + + @classmethod + def get_config_filenames(cls) -> List[str]: + """List of filenames to search for in the model directory.""" + return [ + "quant_config.json", + "quantize_config.json", + ] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> QoQConfig: + weight_bits = cls.get_from_keys(config, ["wbits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + return cls(weight_bits, group_size) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + + if isinstance(layer, LinearBase): + return QoQLinearMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class QoQLinearMethod(LinearMethodBase): + """Linear method for QoQ. + + Args: + quant_config: The QoQ quantization config. + """ + + def __init__(self, quant_config: QoQConfig): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + + weight_loader = extra_weight_attrs.get("weight_loader") + + # Validate output_size_per_partition + output_size_per_partition = sum(output_partition_sizes) + if output_size_per_partition % 32 != 0: + raise ValueError( + f"Weight output_size_per_partition = " + f"{output_size_per_partition} is not divisible by 32." + ) + + # Validate input_size_per_partition + if input_size_per_partition % self.quant_config.pack_factor != 0: + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"pack_factor = {self.quant_config.pack_factor}." + ) + if ( + self.quant_config.group_size != -1 + and input_size_per_partition % self.quant_config.group_size != 0 + ): + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"group_size = {self.quant_config.group_size}." + ) + + qweight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("qweight", qweight) + + s1_scales = ChannelQuantScaleParameter( + data=torch.empty(output_size_per_partition, dtype=torch.float16), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("s1_scales", s1_scales) + + if self.quant_config.group_size == -1: + s1_szeros = ChannelQuantScaleParameter( + data=torch.empty(output_size_per_partition, dtype=torch.float16), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("s1_szeros", s1_szeros) + else: + s2_scales = GroupQuantScaleParameter( + data=torch.empty( + ( + input_size_per_partition // self.quant_config.group_size, + output_size_per_partition, + ), + dtype=torch.int8, + ), + input_dim=0, + output_dim=1, + weight_loader=weight_loader, + ) + layer.register_parameter("s2_scales", s2_scales) + + s2_zeros = GroupQuantScaleParameter( + data=torch.empty( + ( + input_size_per_partition // self.quant_config.group_size, + output_size_per_partition, + ), + dtype=torch.int8, + ), + input_dim=0, + output_dim=1, + weight_loader=weight_loader, + ) + layer.register_parameter("s2_zeros", s2_zeros) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.qweight = Parameter(layer.qweight.data, requires_grad=False) + layer.s1_scales = Parameter(layer.s1_scales.data, requires_grad=False) + if self.quant_config.group_size == -1: + layer.s1_szeros = Parameter(layer.s1_szeros.data, requires_grad=False) + else: + layer.s2_scales = Parameter(layer.s2_scales.data, requires_grad=False) + layer.s2_zeros = Parameter(layer.s2_zeros.data, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + assert x.dtype == torch.float16, "QoQ only supports float16 input now" + if self.quant_config.group_size == -1: + x_q, x_scale, x_sum = per_token_quant_int8( + x, scale_dtype=x.dtype, cal_sum=True + ) + out = qserve_w4a8_per_chn_gemm( + x_q, layer.qweight, layer.s1_scales, x_scale, layer.s1_szeros, x_sum + ) + else: + x_q, x_scale = per_token_quant_int8(x, scale_dtype=x.dtype) + out = qserve_w4a8_per_group_gemm( + x_q, + layer.qweight, + layer.s2_zeros, + layer.s2_scales, + layer.s1_scales, + x_scale, + ) + if bias is not None: + out = out + bias + return out diff --git a/python/sglang/srt/layers/quantization/quark/__init__.py b/python/sglang/srt/layers/quantization/quark/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/layers/quantization/quark/quark.py b/python/sglang/srt/layers/quantization/quark/quark.py new file mode 100644 index 000000000..6d5a66544 --- /dev/null +++ b/python/sglang/srt/layers/quantization/quark/quark.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 + +import fnmatch +import logging +from typing import Any, List, Optional, cast + +import torch + +from sglang.srt.layers.linear import LinearBase, UnquantizedLinearMethod +from sglang.srt.layers.quantization.base_config import ( # noqa: E501 + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod +from sglang.srt.layers.quantization.quark.quark_moe import QuarkMoEMethod +from sglang.srt.layers.quantization.quark.schemes import QuarkScheme, QuarkW4A4MXFP4 +from sglang.srt.layers.quantization.quark.utils import deep_compare, should_ignore_layer +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.utils import get_device_capability + +__all__ = ["QuarkLinearMethod"] + +logger = logging.getLogger(__name__) + + +class QuarkConfig(QuantizationConfig): + + def __init__( + self, + quant_config: dict[str, Any], + kv_cache_group: Optional[list[str]] = None, + kv_cache_config: Optional[dict[str, Any]] = None, + pack_method: str = "reorder", + ): + super().__init__() + if kv_cache_group is None: + kv_cache_group = [] + self.quant_config = quant_config + self.kv_cache_group = kv_cache_group + self.kv_cache_config = kv_cache_config + self.pack_method = pack_method + + self.packed_modules_mapping = self.quant_config["packed_modules_mapping"] + + def get_linear_method(self) -> "QuarkLinearMethod": + return QuarkLinearMethod(self) + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.float16, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 70 + + def get_name(self) -> str: + return "quark" + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["QuantizeMethodBase"]: + # Check if the layer is skipped for quantization. + exclude_layers = cast(list[str], self.quant_config.get("exclude")) + if should_ignore_layer( + prefix, ignore=exclude_layers, fused_mapping=self.packed_modules_mapping + ): + return UnquantizedLinearMethod() + + if isinstance(layer, LinearBase): + scheme = self.get_scheme(layer=layer, layer_name=prefix) + layer.scheme = scheme + return QuarkLinearMethod(self) + + if isinstance(layer, RadixAttention): + return QuarkKVCacheMethod(self) + + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + + if isinstance(layer, FusedMoE): + return QuarkMoEMethod.get_moe_method(self, module=layer, layer_name=prefix) + + return None + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "QuarkConfig": + export_config = config.get("export") + if export_config is None: + raise ValueError( + "The export key should be included in " + "the configurations of Quark quantized model" + ) + + kv_cache_group = cast(list[str], export_config.get("kv_cache_group")) + pack_method = cast(str, export_config.get("pack_method")) + + # In the export model of quark, the quantization configuration + # of kv_cache is stored in layer_quant_config. First, it is + # judged whether kv_cache_group exists, and then it is judged + # whether layer_quant_config has a quantization configuration + # that matches kv_cache. + if len(kv_cache_group) == 0: + kv_cache_config = None + else: + kv_cache_set = set(kv_cache_group) + layer_quant_config = cast(dict[str, Any], config.get("layer_quant_config")) + layer_quant_names = list(layer_quant_config.keys()) + layer_quant_set = set(layer_quant_names) + + if not kv_cache_set.issubset(layer_quant_set): + raise ValueError( + "The Quark quantized model has the " + "kv_cache_group parameter setting, " + "but no kv_cache quantization settings " + "were found in the quantization " + "configuration." + ) + + q_configs = [ + cast(dict[str, Any], layer_quant_config.get(name)) + for name in kv_cache_group + ] + if not all(deep_compare(q_config, q_configs[0]) for q_config in q_configs): + raise ValueError( + "The quantization method used for kv_cache should " + "be the same, but the quantization method for the " + "kv_cache layer in the config is different." + ) + kv_cache_config = q_configs[0].get("output_tensors") + if kv_cache_config is None: + raise ValueError("The kv_cache quantization configuration is empty.") + + # Since we have already set kv_cache quantization configurations, + # we will remove the quantization configuration for the + # output_tensors corresponding to the kv_cache layer. + for q_config in q_configs: + q_config["output_tensors"] = None + + # In case q_proj output is also quantized, remove the configuration + # to keep qkv consistency. + q_proj_q_config = cast(dict[str, Any], layer_quant_config.get("*q_proj")) + if q_proj_q_config is not None: + q_proj_q_config["output_tensors"] = None + + return cls( + quant_config=config, + kv_cache_group=kv_cache_group, + kv_cache_config=kv_cache_config, + pack_method=pack_method, + ) + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def _check_scheme_supported(self, min_capability: int, error: bool = True) -> bool: + capability_tuple = get_device_capability() + + if capability_tuple is not None: + assert 0 <= capability_tuple[1] < 10 + capability = capability_tuple[0] * 10 + capability_tuple[1] + + supported = capability >= min_capability + if error and not supported: + raise RuntimeError( + "Quantization scheme is not supported for ", + f"the current GPU. Min capability: {min_capability}. ", + f"Current capability: {capability}.", + ) + return supported + else: + return False + + def _is_mx_fp4( + self, + weight_quant: Optional[dict[str, Any]], + input_quant: Optional[dict[str, Any]], + ) -> bool: + # Confirm weights and input quantized. + if weight_quant is None or input_quant is None: + logger.debug( + "Quark model is not in MX-FP4 format: " + "weight_quant or input_quant not set" + ) + return False + + # Input and weight dtype needs to be fp4. + if weight_quant.get("dtype") != "fp4" or input_quant.get("dtype") != "fp4": + logger.debug("Quark model is not in MX-FP4 format: dtype not fp4") + return False + + # Input and weight qscheme needs to be per group. + if ( + weight_quant.get("qscheme") != "per_group" + or input_quant.get("qscheme") != "per_group" + ): + logger.debug("Quark model is not in MX-FP4 format: not per_group") + return False + + # Input and weight group size needs to be 32. + if weight_quant.get("group_size") != 32 or input_quant.get("group_size") != 32: + logger.debug("Quark model is not in MX-FP4 format: not group_size=32") + return False + + # Weights need to use static quantization. + if weight_quant.get("is_dynamic") is True: + logger.debug("Quark model is not in MX-FP4 format: not weight static") + return False + + # Activations need to use dynamic quantization. + if input_quant.get("is_dynamic") is False: + logger.debug("Quark model is not in MX-FP4 format: not activation dynamic") + return False + + # Activations and weight scales need to be in e8m0 format. + if ( + weight_quant.get("scale_format") != "e8m0" + or input_quant.get("scale_format") != "e8m0" + ): + logger.debug("Quark model is not in MX-FP4 format: not scale_format e8m0") + return False + + return True + + def _find_matched_config( + self, layer_name: str, module: torch.nn.Module + ) -> dict[str, Any]: + + proj_name = layer_name.split(".")[-1] + if proj_name in self.packed_modules_mapping: + shard_proj_names = self.packed_modules_mapping[proj_name] + + # Convert fused_name --> [shard_names] + shard_names = [ + layer_name.replace(proj_name, shard_proj_name) + for shard_proj_name in shard_proj_names + ] + shard_configs = [ + self._find_matched_config(shard_name, module) + for shard_name in shard_names + ] + if not all( + deep_compare(q_config, shard_configs[0]) for q_config in shard_configs + ): + raise ValueError( + f"Found a different quantization configuration for " + f"{shard_proj_names} in {layer_name}. vLLM " + "requires all to use the same scheme." + ) + return shard_configs[0] + else: + layer_quant_config = cast( + dict[str, Any], self.quant_config.get("layer_quant_config") + ) + for name_pattern in layer_quant_config: + if fnmatch.fnmatch(layer_name, name_pattern): + return layer_quant_config[name_pattern] + + layer_type = type(module).__name__ + layer_type_quant_config = cast( + dict[str, Any], self.quant_config.get("layer_type_quant_config") + ) + if layer_type in layer_type_quant_config: + return layer_type_quant_config[layer_type] + + global_quant_config = cast( + dict[str, Any], self.quant_config.get("global_quant_config") + ) + return global_quant_config + + def _get_scheme_from_config(self, config: dict[str, Any]) -> "QuarkScheme": + if config.get("output_tensors") or config.get("bias"): + raise NotImplementedError( + "Currently, Quark models with output_tensors " + "and bias quantized are not supported" + ) + weight_config = cast(dict[str, Any], config.get("weight")) + input_config = cast(dict[str, Any], config.get("input_tensors")) + + if self._is_mx_fp4(weight_config, input_config): + return QuarkW4A4MXFP4(weight_config, input_config) + + raise NotImplementedError( + "No quark compatible scheme was found. " + f"Weight config: {weight_config}, " + f"Input config: {input_config}" + ) + + def get_scheme(self, layer: torch.nn.Module, layer_name: str) -> "QuarkScheme": + + layer_quant_config = self._find_matched_config(layer_name, layer) + + # Find the quant_scheme + scheme = self._get_scheme_from_config(layer_quant_config) + + # Raise error if device does not support the scheme + # (e.g. fp8 needs ada lovelace) + self._check_scheme_supported(scheme.get_min_capability()) + + return scheme + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class QuarkLinearMethod(LinearMethodBase): + + def __init__(self, quantization_config: QuarkConfig): + self.quantization_config = quantization_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + """ + Use the CompressedTensorsScheme associated with each layer to create + the necessary parameters for the layer. See LinearMethodBase for param + details + """ + weight_loader = extra_weight_attrs.get("weight_loader") + layer.scheme.create_weights( + layer=layer, + input_size=input_size, + input_size_per_partition=input_size_per_partition, + output_partition_sizes=output_partition_sizes, + output_size=output_size, + params_dtype=params_dtype, + weight_loader=weight_loader, + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + """ + Use the output of create_weights and the CompressedTensorsScheme + associated with the layer to apply the forward pass with the + layer input. See LinearMethodBase for param details + + """ + scheme = layer.scheme + if scheme is None: + raise ValueError("A scheme must be defined for each layer") + return scheme.apply_weights(layer, x, bias=bias) + + +class QuarkKVCacheMethod(BaseKVCacheMethod): + """ + Supports loading kv-cache scaling factors from quark checkpoints. + """ + + def __init__(self, quant_config: QuarkConfig): + self.validate_kv_cache_config(quant_config.kv_cache_config) + super().__init__(quant_config) + + @staticmethod + def validate_kv_cache_config(kv_cache_config: Optional[dict[str, Any]]): + """ + Validator for the kv cache configuration. Useful for controlling the + kv cache quantization schemes, that are being supported in vLLM + :param kv_cache_config: the quark kv cache scheme + """ + if kv_cache_config is None: + return + + dtype = kv_cache_config.get("dtype") + if dtype != "fp8_e4m3": + raise NotImplementedError( + "Currently supported kv cache quantization is " + f"dtype=fp8_e4m3, however received {dtype}" + ) + + qscheme = kv_cache_config.get("qscheme") + if qscheme != "per_tensor": + raise NotImplementedError( + "Only support per-tensor scaling factor " + "for quark KV cache. " + f"Expected qscheme: per_tensor, found qscheme: {qscheme}" + ) diff --git a/python/sglang/srt/layers/quantization/quark/quark_moe.py b/python/sglang/srt/layers/quantization/quark/quark_moe.py new file mode 100644 index 000000000..f6e750a2c --- /dev/null +++ b/python/sglang/srt/layers/quantization/quark/quark_moe.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Callable, Optional + +import torch +from aiter import ActivationType, QuantType, biased_grouped_topk +from aiter.fused_moe import fused_moe +from aiter.utility.fp4_utils import e8m0_shuffle + +from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase +from sglang.srt.utils import get_bool_env_var, mxfp_supported, set_weight_attrs + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + from sglang.srt.layers.quantization.quark.quark import QuarkConfig + +logger = logging.getLogger(__name__) + +__all__ = ["QuarkMoEMethod", "QuarkW4A4MXFp4MoEMethod"] + +OCP_MX_BLOCK_SIZE = 32 + +if TYPE_CHECKING: + from sglang.srt.layers.quantization import QuarkConfig + + +class QuarkMoEMethod(FusedMoEMethodBase): + + def __init__(self, quant_config: QuarkConfig): + self.quant_config = quant_config + + @staticmethod + def get_moe_method( + quant_config: QuarkConfig, # type: ignore # noqa E501 # noqa F821 + module: torch.nn.Module, + layer_name: str, + ) -> "QuarkMoEMethod": + layer_quant_config = quant_config._find_matched_config(layer_name, module) + + if layer_quant_config.get("output_tensors") or layer_quant_config.get("bias"): + raise NotImplementedError( + "Currently, Quark models with " + "output_tensors and bias " + "quantized are not supported" + ) + weight_config = layer_quant_config.get("weight") + input_config = layer_quant_config.get("input_tensors") + + if quant_config._is_mx_fp4(weight_config, input_config): + return QuarkW4A4MXFp4MoEMethod(weight_config, input_config) + else: + raise RuntimeError("Unsupported FusedMoe scheme") + + +class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod): + + def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]): + self.weight_quant = weight_config + self.input_quant = input_config + + weight_qscheme = self.weight_quant.get("qscheme") + input_qscheme = self.input_quant.get("qscheme") + if not (weight_qscheme == "per_group" and input_qscheme == "per_group"): + raise ValueError( + "For MX(FP4) Fused MoE layers, only per-group scales " + "for weights and activations are supported. Found " + f"{weight_qscheme}, {input_qscheme}" + ) # noqa E501 + + self.static_input_scales = not self.input_quant.get("is_dynamic") + self.with_bias = False + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + # Add the quantization method used (per tensor/grouped/channel) + # to ensure the weight scales are loaded in properly + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + ) + + params_dtype = torch.uint8 + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // 2, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // 2, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // OCP_MX_BLOCK_SIZE, + dtype=params_dtype, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + hidden_size, + intermediate_size_per_partition // OCP_MX_BLOCK_SIZE, + dtype=params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + float_dtype = torch.get_default_dtype() + + # Pre-shuffle weight scales + s0, s1, _ = layer.w13_weight_scale.shape + w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1) + w13_weight_scale = e8m0_shuffle(w13_weight_scale) + # layer.w13_weight_scale = torch.nn.Parameter(w13_weight_scale, requires_grad=False) + layer.w13_weight_scale.data = w13_weight_scale.view(s0, s1, -1) + + s0, s1, _ = layer.w2_weight_scale.shape + w2_weight_scale = layer.w2_weight_scale.view(s0 * s1, -1) + w2_weight_scale = e8m0_shuffle(w2_weight_scale) + # layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False) + layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + moe_runner_config = self.moe_runner_config + topk_weights, topk_ids, _ = topk_output + + output = fused_moe( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + quant_type=QuantType.per_1x32, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + activation=( + ActivationType.Silu + if moe_runner_config.activation == "silu" + else ActivationType.Gelu + ), + doweight_stage1=False, + ) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/quark/schemes/__init__.py b/python/sglang/srt/layers/quantization/quark/schemes/__init__.py new file mode 100644 index 000000000..91b08c512 --- /dev/null +++ b/python/sglang/srt/layers/quantization/quark/schemes/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 + +from .quark_scheme import QuarkScheme +from .quark_w4a4_mxfp4 import QuarkW4A4MXFP4 + +__all__ = ["QuarkScheme", "QuarkW4A4MXFP4"] diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_scheme.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_scheme.py new file mode 100644 index 000000000..aab5c9c1e --- /dev/null +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_scheme.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod +from typing import Optional + +import torch + +__all__ = ["QuarkScheme"] + + +class QuarkScheme(ABC): + """ + Abstract class used to describe the weight creation and forward pass + of different quantization schemes supported by Quark. + """ + + @classmethod + @abstractmethod + def get_min_capability(cls) -> int: + """ + Get minimum device capability. + """ + raise NotImplementedError + + @abstractmethod + def create_weights(self, *args, **kwargs): + """ + Weight creation for the particular scheme. Inputs to this function + + """ + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] + ): + """ + Run the forward pass for the particular scheme. This is where + scheme-specific dequant/quant steps/kernels should be applied. + + :param layer: torch.nn.Module with the registered weights and + other parameters relevant to the particular scheme. + :param x: input to the layer + :param bias: bias parameter + + """ + raise NotImplementedError + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module): + """ + Called after weight loading is complete for any cleanup that + needs to occur. + """ + raise NotImplementedError diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py new file mode 100644 index 000000000..a0787baaf --- /dev/null +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Callable, Optional + +import aiter +import torch +import torch.nn.functional as F +from aiter.ops.gemm_op_a4w4 import gemm_a4w4 +from aiter.ops.shuffle import shuffle_weight +from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4 +from aiter.ops.triton.gemm_afp4wfp4_pre_quant_atomic import gemm_afp4wfp4_pre_quant +from aiter.ops.triton.quant import dynamic_mxfp4_quant +from aiter.utility import dtypes +from aiter.utility.fp4_utils import e8m0_shuffle + +from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter +from sglang.srt.layers.quantization.quark.schemes import QuarkScheme +from sglang.srt.utils import get_bool_env_var + +__all__ = ["QuarkW4A4MXFP4"] + +OCP_MX_BLOCK_SIZE = 32 + + +class QuarkW4A4MXFP4(QuarkScheme): + + def __init__( + self, weight_quant_spec: dict[str, Any], input_quant_spec: dict[str, Any] + ): + self.out_dtype = torch.get_default_dtype() + self.qscheme = "per_group" + self.weight_quant_spec = weight_quant_spec + self.input_quant_spec = input_quant_spec + + @classmethod + def get_min_capability(cls) -> int: + return 70 + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + return + + def create_weights( + self, + layer: torch.nn.Module, + output_partition_sizes: list[int], + input_size_per_partition: int, + params_dtype: torch.dtype, + weight_loader: Callable, + **kwargs + ): + output_size_per_partition = sum(output_partition_sizes) + layer.logical_widths = output_partition_sizes + + # WEIGHT + weight = PackedvLLMParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + packed_dim=1, + packed_factor=2, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + weight_scale = GroupQuantScaleParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // OCP_MX_BLOCK_SIZE, + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # This path does not have support for bias currently + assert bias is None, "bias is not supported" + + three_d = False + x_s = None + y = None + if isinstance(x, tuple): + assert len(x) in [ + 2, + 3, + ], "For tuple input, only (x, x_s) or (x, x_s, y) formats are accepted" + if len(x) == 2: + x, x_s = x + elif len(x) == 3: + x, x_s, y = x + + use_fused_quant_gemm = ( + x_s is None and y is not None and layer.weight.shape[0] == y.shape[1] + ) + + if x.dim() == 3: + three_d = True + x = x.view(-1, x.shape[-1]) + output_shape = [*x.shape[:-1], layer.weight.shape[0]] + + # use_fused_quant_gemm = true, x_q is a bf16/fp16 num + # x_s is not None = true, x_q is uint8 num + if use_fused_quant_gemm or x_s is not None: + x_q = x + else: + x_q, x_s = dynamic_mxfp4_quant(x) + + if y is None: + y = torch.empty( + x_q.shape[0], + layer.weight.shape[0], + device=x_q.device, + dtype=self.out_dtype, + ) + + if use_fused_quant_gemm: + gemm_afp4wfp4_pre_quant(x_q, layer.weight, layer.weight_scale, y.dtype, y) + y = y.to(x.dtype) + else: + gemm_afp4wfp4(x_q, layer.weight, x_s, layer.weight_scale, self.out_dtype, y) + + if three_d: + return y.view(*output_shape) + + return y diff --git a/python/sglang/srt/layers/quantization/quark/utils.py b/python/sglang/srt/layers/quantization/quark/utils.py new file mode 100644 index 000000000..eacbf3ba9 --- /dev/null +++ b/python/sglang/srt/layers/quantization/quark/utils.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 + +import re +from collections.abc import Iterable, Mapping +from types import MappingProxyType +from typing import Any, Optional + +import torch +from aiter.ops.triton.quant import dynamic_mxfp4_quant +from torch import nn + + +def deep_compare(dict1: Any, dict2: Any) -> bool: + if type(dict1) is not type(dict2): + return False + if isinstance(dict1, dict): + if dict1.keys() != dict2.keys(): + return False + return all(deep_compare(dict1[k], dict2[k]) for k in dict1) + elif isinstance(dict1, list): + return set(dict1) == set(dict2) + else: + return dict1 == dict2 + + +def should_ignore_layer( + layer_name: Optional[str], + ignore: Iterable[str], + fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), +) -> bool: + if layer_name is None: + return False + + # layer_name = model.layers.0.self_attn.qkv_proj + # proj_name = qkv_proj + proj_name = layer_name.split(".")[-1] + + # Fused layers like gate_up_proj or qkv_proj will not be fused + # in the safetensors checkpoint. So, we convert the name + # from the fused version to unfused + check to make sure that + # each shard of the fused layer has the same scheme. + if proj_name in fused_mapping: + shard_proj_names = fused_mapping[proj_name] + + # Convert fused_name --> [shard_names] + shard_names = [ + layer_name.replace(proj_name, shard_proj_name) + for shard_proj_name in shard_proj_names + ] + + # Layer should be ignored if shards are ignored. + should_ignore_layer = None + for shard_name in shard_names: + should_ignore_shard = check_equal_or_regex_match( + layer_name=shard_name, targets=ignore + ) + + # If shard_idx=0, set layer ignore to match shard. + if should_ignore_layer is None: + should_ignore_layer = should_ignore_shard + + # If shard_idx=1+ confirm scheme matches prior shards. + elif should_ignore_shard != should_ignore_layer: + raise ValueError( + f"Found a different quantization schemes for " + f"{shard_proj_names} in {layer_name}. vLLM " + "requires all to use the same scheme." + ) + + # Unfused layers like down_proj and o_proj will match + # the safetensors checkpoint already. + else: + should_ignore_layer = check_equal_or_regex_match( + layer_name=layer_name, targets=ignore + ) + + assert should_ignore_layer is not None + + return should_ignore_layer + + +def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool: + """ + Checks whether a layer_name is exactly equal or a regex match for + if target starts with 're:' to any target in list. + """ + for target in targets: + if _is_equal_or_regex_match(layer_name, target): + return True + return False + + +def _is_equal_or_regex_match( + value: str, target: str, check_contains: bool = False +) -> bool: + """ + Checks whether a value is exactly equal or a regex match for target + if target starts with 're:'. If check_contains is set to True, + additionally checks if the target string is contained within the value. + """ + + if target.startswith("re:"): + pattern = target[3:] + if re.match(pattern, value): + return True + elif check_contains: + if target.lower() in value.lower(): + return True + elif target == value: + return True + return False + + +# utility for tensor dims > 2 cases +def b_dynamic_mxfp4_quant(x): + h, b, d = x.shape + x, x_scales = dynamic_mxfp4_quant(x.reshape(-1, d)) + return x.view(h, b, d // 2), x_scales.view(h, b, d // 32) + + +def mxfp4_to_f32(x, is_threed): + # 2 because we pack fp4 in uint8. + x = x.repeat_interleave(2, dim=-1) + if is_threed: + x[..., ::2] = x[..., ::2] & 0xF + x[..., 1::2] = x[..., 1::2] >> 4 + else: + x[:, ::2] = x[:, ::2] & 0xF + x[:, 1::2] = x[:, 1::2] >> 4 + + mxfp4_list = [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ] + mxfp4_in_f32 = torch.tensor(mxfp4_list, dtype=torch.float32, device="cuda") + return mxfp4_in_f32[x.long()] + + +def e8m0_to_f32(x): + # Convert the input tensor `x` (assumed to be in e8m0 format) to float32. + # e8m0 is a custom 8-bit floating point format with 8 bits for exponent, 0 for mantissa. + # This means the value is essentially 2^(exponent - 127), similar to how IEEE-754 stores floats. + + # Convert x to float32 for computation, and compute the power of 2 by subtracting the bias (127). + x_f32 = 2 ** ((x.to(torch.float32)) - 127) + + # If the exponent value was 255 (i.e., 2^(128)), this is a special case usually used to represent NaN or Inf. + # Since this custom format has no mantissa, treat 2^128 as NaN. + x_f32[x_f32 == 128] = float("nan") + return x_f32 + + +def quark_post_load_weights(self_attn: nn.Module, w: torch.Tensor, quant_format: str): + if "mxfp4" in quant_format: + # when dtype is bf16, the processing flow is to dynamic quantize bf16 tensor to uint8 tensor + # do w_kc (bf16) first to get the w_kc(uint8) w_s_kc(uint8) + # and w_vc repeating the same procedure of w_kc to get w_vc(uint8) w_s_vc(uint8) + if w.dtype == torch.bfloat16: + w_kc, w_vc = w.unflatten( + 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) + ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) + w_kc, w_s_kc = b_dynamic_mxfp4_quant(w_kc.transpose(-2, -1)) + w_kc = w_kc.transpose(-2, -1) + w_s_kc = w_s_kc.transpose(-2, -1) + w_vc, w_s_vc = b_dynamic_mxfp4_quant(w_vc) + w_s_kc = w_s_kc.transpose(1, 2).contiguous().transpose(1, 2) + w_s_vc = w_s_vc.contiguous().transpose(1, 2) + elif w.dtype == torch.uint8: # static quant for mxfp4 + # when dtype is uint8, it means the w has been quantized to mxfp4 format + # but we must separate it to w_kc and w_vc. + # The quantized tensor size is only half of original tensor size + # and the scaling factor is 1/32, the transpose behavior will be not correct + # need to upcast it to fp32 to separate w to w_kc and w_vc + # to ensure the following transpose behavior is correct + # and then do mxfp4 quant again + w = mxfp4_to_f32(w, True).to(torch.bfloat16) + w_scales = self_attn.kv_b_proj.weight_scale.repeat_interleave(32, dim=-1) + w_scales = e8m0_to_f32(w_scales).to(torch.bfloat16) + w = w * w_scales + w_kc, w_vc = w.unflatten( + 0, (-1, (self_attn.qk_nope_head_dim + self_attn.v_head_dim)) + ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) + w_kc, w_s_kc = b_dynamic_mxfp4_quant(w_kc.transpose(-2, -1)) + w_kc = w_kc.transpose(-2, -1) + w_s_kc = w_s_kc.transpose(-2, -1) + w_vc, w_s_vc = b_dynamic_mxfp4_quant(w_vc) + w_s_kc = w_s_kc.transpose(1, 2).contiguous().transpose(1, 2) + w_s_vc = w_s_vc.contiguous().transpose(1, 2) + + return w_kc, w_s_kc, w_vc, w_s_vc diff --git a/python/sglang/srt/layers/quantization/rocm_mxfp4_utils.py b/python/sglang/srt/layers/quantization/rocm_mxfp4_utils.py new file mode 100644 index 000000000..4659f76bd --- /dev/null +++ b/python/sglang/srt/layers/quantization/rocm_mxfp4_utils.py @@ -0,0 +1,13 @@ +from aiter.ops.triton.batched_gemm_afp4wfp4_pre_quant import ( + batched_gemm_afp4wfp4_pre_quant, +) +from aiter.ops.triton.fused_mxfp4_quant import ( + fused_flatten_mxfp4_quant, + fused_rms_mxfp4_quant, +) + +__all__ = [ + "fused_rms_mxfp4_quant", + "fused_flatten_mxfp4_quant", + "batched_gemm_afp4wfp4_pre_quant", +] diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py new file mode 100644 index 000000000..495beb009 --- /dev/null +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -0,0 +1,468 @@ +from __future__ import annotations + +import importlib.util +from typing import TYPE_CHECKING, List, Optional + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from sglang.srt.custom_op import CustomOp +from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizeMethodBase, +) +from sglang.srt.utils import ( + cpu_has_amx_support, + get_bool_env_var, + is_cpu, + is_hip, + set_weight_attrs, + use_intel_amx_backend, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +has_triton_kernels = importlib.util.find_spec("triton_kernels") is not None + + +_is_cpu_amx_available = cpu_has_amx_support() +_is_hip = is_hip() +_is_cpu = is_cpu() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _use_aiter: + from aiter import ActivationType + from aiter.fused_moe import fused_moe + from aiter.ops.shuffle import shuffle_weight + + +class UnquantizedEmbeddingMethod(QuantizeMethodBase): + """Unquantized method for embeddings.""" + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + """Create weights for embedding layer.""" + weight = Parameter( + torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return F.linear(x, layer.weight, bias) + + def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor: + return F.embedding(input_, layer.weight) + + +class UnquantizedLinearMethod(LinearMethodBase): + """Linear method without quantization.""" + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + weight = Parameter( + torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if _is_cpu and _is_cpu_amx_available: + _amx_process_weight_after_loading(layer, ["weight"]) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + if use_intel_amx_backend(layer): + x_shapes = x.shape + if len(x_shapes) == 3: + x = x.view(-1, x.shape[-1]) + output = torch.ops.sgl_kernel.weight_packed_linear( + x, layer.weight, bias, True # is_vnni + ) + if len(x_shapes) == 3: + output = output.view(x_shapes[0], x_shapes[1], -1) + return output + + return F.linear(x, layer.weight, bias) + + +class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): + """MoE method without quantization.""" + + def __init__(self, use_triton_kernels: bool = False): + super().__init__() + self.use_triton_kernels = use_triton_kernels + self.with_bias = False + + self.triton_kernel_moe_forward = None + self.triton_kernel_moe_with_bias_forward = None + if torch.cuda.is_available() and has_triton_kernels: + from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import ( + triton_kernel_moe_forward as _tk_forward, + ) + from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import ( + triton_kernel_moe_with_bias_forward as _tk_with_bias_forward, + ) + + self.triton_kernel_moe_forward = _tk_forward + self.triton_kernel_moe_with_bias_forward = _tk_with_bias_forward + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + with_bias: bool = False, + **extra_weight_attrs, + ): + self.with_bias = with_bias + + # Fused gate_up_proj (column parallel) + w13_weight_n, w13_weight_k = 2 * intermediate_size_per_partition, hidden_size + if self.use_triton_kernels: + w13_weight_n, w13_weight_k = w13_weight_k, w13_weight_n + w13_weight = torch.nn.Parameter( + torch.empty(num_experts, w13_weight_n, w13_weight_k, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + if self.with_bias: + w13_weight_bias = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_bias", w13_weight_bias) + set_weight_attrs(w13_weight_bias, extra_weight_attrs) + + # down_proj (row parallel) + w2_weight_n, w2_weight_k = ( + hidden_size, + intermediate_size_per_partition, + ) + if self.use_triton_kernels: + w2_weight_n, w2_weight_k = w2_weight_k, w2_weight_n + w2_weight = torch.nn.Parameter( + torch.empty(num_experts, w2_weight_n, w2_weight_k, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + if self.with_bias: + w2_weight_bias = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_bias", w2_weight_bias) + set_weight_attrs(w2_weight_bias, extra_weight_attrs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if _use_aiter: + layer.w13_weight = torch.nn.Parameter( + shuffle_weight(layer.w13_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + layer.w2_weight = torch.nn.Parameter( + shuffle_weight(layer.w2_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + + # Pack weight for get better performance on CPU + if _is_cpu and _is_cpu_amx_available: + _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"]) + + return + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + return self.forward( + layer=layer, + dispatch_output=dispatch_output, + ) + + def forward_cuda( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + moe_runner_config = self.moe_runner_config + + if self.use_triton_kernels: + if self.with_bias: + assert self.triton_kernel_moe_with_bias_forward is not None + output = self.triton_kernel_moe_with_bias_forward( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + b1=layer.w13_weight_bias, + b2=layer.w2_weight_bias, + topk_output=topk_output, + moe_runner_config=moe_runner_config, + w1_pcg=None, + w2_pcg=None, + ) + else: + assert self.triton_kernel_moe_forward is not None + output = self.triton_kernel_moe_forward( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_output=topk_output, + moe_runner_config=moe_runner_config, + ) + return StandardCombineInput(hidden_states=output) + else: + if _use_aiter: + assert not moe_runner_config.no_combine, "unsupported" + topk_weights, topk_ids, _ = topk_output + if moe_runner_config.apply_router_weight_on_input: + assert ( + topk_weights.dim() == 2 + ), "`topk_weights` should be in shape (num_tokens, topk)" + _, topk = topk_weights.shape + assert ( + topk == 1 + ), "Only support topk=1 when `apply_router_weight_on_input` is True" + x = x * topk_weights.to(x.dtype) + topk_weights = torch.ones_like( + topk_weights, dtype=torch.float32 + ) # topk_weights must be FP32 (float32) + output = fused_moe( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=( + ActivationType.Silu + if moe_runner_config.activation == "silu" + else ActivationType.Gelu + ), + ) + return StandardCombineInput(hidden_states=output) + else: + + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + b13=getattr(layer, "w13_weight_bias", None), + b2=getattr(layer, "w2_weight_bias", None), + ) + return self.runner.run(dispatch_output, quant_info) + + def forward_cpu( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + moe_runner_config = self.moe_runner_config + + assert ( + moe_runner_config.activation == "silu" + ), f"activation = {moe_runner_config.activation} is not supported." + + if ( + use_intel_amx_backend(layer) + and not moe_runner_config.apply_router_weight_on_input + ): + from sglang.srt.layers.moe.topk import apply_topk_weights_cpu + + topk_weights, topk_ids, _ = topk_output + x, topk_weights = apply_topk_weights_cpu( + moe_runner_config.apply_router_weight_on_input, topk_weights, x + ) + output = torch.ops.sgl_kernel.fused_experts_cpu( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + False, # inplace # See [Note] inplace should be False in fused_experts. + False, # use_int8_w8a8 + False, # use_fp8_w8a16 + None, # w1_scale + None, # w2_scale + None, # block_size + None, # a1_scale + None, # a2_scale + True, # is_vnni + ) + return StandardCombineInput(hidden_states=output) + else: + from sglang.srt.layers.moe.fused_moe_native import moe_forward_native + + output = moe_forward_native( + layer, + x, + topk_output, + moe_runner_config, + ) + return StandardCombineInput(hidden_states=output) + + def forward_npu( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + import torch_npu + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_weights, topk_ids, _ = dispatch_output.topk_output + + original_dtype = x.dtype + num_tokens = x.shape[0] + topk_weights = topk_weights.to(x.dtype) + topk_ids = topk_ids.to(torch.int32) + num_experts = layer.num_experts + top_k = layer.top_k + row_idx_len = num_tokens * top_k + row_idx = ( + torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device) + .view(top_k, -1) + .permute(1, 0) + .contiguous() + ) + + hidden_states, expanded_row_idx, expanded_expert_idx = ( + torch_npu.npu_moe_init_routing( + x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens + ) + ) + + expert_tokens = torch_npu.npu_moe_compute_expert_tokens( + expanded_expert_idx, num_experts + ) + + expert_tokens = expert_tokens.to(torch.int64) + if layer.w13_weight.shape[-1] == layer.hidden_size: + w13 = layer.w13_weight.transpose(1, 2) + w2 = layer.w2_weight.transpose(1, 2) + + # gmm1: gate_up_proj + hidden_states = torch_npu.npu_grouped_matmul( + x=[hidden_states], + weight=[w13], + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + + # act_fn: + if self.moe_runner_config.activation == "silu": + hidden_states = torch_npu.npu_swiglu(hidden_states) + else: + from sglang.srt.layers.activation import GeluAndMul + + hidden_states = GeluAndMul()(hidden_states) + + # gmm2: down_proj + hidden_states = torch_npu.npu_grouped_matmul( + x=[hidden_states], + weight=[w2], + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + + final_hidden_states = torch_npu.npu_moe_finalize_routing( + hidden_states, + skip1=None, + skip2=None, + bias=None, + scales=topk_weights, + expanded_src_to_dst_row=expanded_row_idx, + export_for_source_row=topk_ids, + ) + + return StandardCombineInput(hidden_states=final_hidden_states) + + def forward_tpu(self, *args, **kwargs) -> CombineInput: + raise NotImplementedError("The TPU backend currently does not support MoE.") + + forward_native = forward_cpu diff --git a/python/sglang/srt/layers/quantization/utils.py b/python/sglang/srt/layers/quantization/utils.py new file mode 100644 index 000000000..63b8b6eb7 --- /dev/null +++ b/python/sglang/srt/layers/quantization/utils.py @@ -0,0 +1,564 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/quant_utils.py + +from __future__ import annotations + +import re +from copy import deepcopy +from types import MappingProxyType +from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Tuple, Union + +import numpy +import torch + +from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant +from sglang.srt.utils import is_cuda + +if TYPE_CHECKING: + from sglang.srt.layers.quantization.base_config import QuantizationConfig + + +def get_scalar_types(): + """ + Returns: + tuple: (ScalarType, scalar_types) + """ + try: + from sgl_kernel.scalar_type import ScalarType, scalar_types + + return ScalarType, scalar_types + except ImportError: + + class MockScalarType: + pass + + class MockScalarTypes: + uint4b8 = "uint4b8" + uint8b128 = "uint8b128" + + def __getattr__(self, name): + return f"mock_{name}" + + return MockScalarType, MockScalarTypes() + + +ScalarType, scalar_types = get_scalar_types() + + +def is_layer_skipped( + prefix: str, + ignored_layers: List[str], + fused_mapping: Mapping[str, List[str]] = MappingProxyType({}), +) -> bool: + # prefix: model.layers.0.self_attn.q_proj + # proj_name: q_proj + proj_name = prefix.split(".")[-1] + + # Fused layers like gate_up_proj or qkv_proj will not be fused + # in the safetensors checkpoint. So, we convert the name + # from the fused version to unfused + check to make sure that + # each shard of the fused layer has the same scheme. + if proj_name in fused_mapping: + shard_prefixes = [ + prefix.replace(proj_name, shard_proj_name) + for shard_proj_name in fused_mapping[proj_name] + ] + + is_skipped = None + for shard_prefix in shard_prefixes: + is_shard_skipped = shard_prefix in ignored_layers + + if is_skipped is None: + is_skipped = is_shard_skipped + elif is_shard_skipped != is_skipped: + raise ValueError( + f"Detected some but not all shards of {prefix} " + "are quantized. All shards of fused layers " + "to have the same precision." + ) + else: + is_skipped = prefix in ignored_layers + if "gate_up_proj" in prefix: + prefix_gate = prefix.replace("gate_up_proj", "gate_proj") + prefix_up = prefix.replace("gate_up_proj", "up_proj") + if prefix_gate in ignored_layers and prefix_up in ignored_layers: + is_skipped = True + elif "experts" in prefix: + is_skipped = any( + [ + prefix in layer_name + for layer_name in ignored_layers + if "experts" in layer_name + ] + ) + + assert is_skipped is not None + return is_skipped + + +def per_tensor_dequantize( + tensor: torch.Tensor, inv_scale: Union[float, torch.Tensor] +) -> torch.Tensor: + fake_qweight = tensor.to(torch.float16) + dq_weight = fake_qweight * inv_scale + return dq_weight + + +def all_close_1d(x: torch.Tensor) -> bool: + assert len(x.shape) == 1 + return all(torch.allclose(x[0], x[i]) for i in range(x.shape[0])) + + +def convert_to_channelwise( + weight_scale: torch.Tensor, logical_widths: List[int] +) -> Tuple[torch.Tensor, torch.Tensor]: + # Create channelwise buffer + weight_scale_channel = torch.empty( + (sum(logical_widths), 1), dtype=torch.float32, device=weight_scale.device + ) + + # Handle scalar tensor case: broadcast same scale to all channels + if weight_scale.dim() == 0: + weight_scale_channel.fill_(weight_scale.item()) + return weight_scale_channel + + # Expand each scale to match the size of each logical matrix. + start = 0 + for idx, logical_width in enumerate(logical_widths): + end = start + logical_width + weight_scale_channel[start:end, :] = weight_scale[idx] + start = end + + return weight_scale_channel + + +def requantize_with_max_scale( + weight: torch.Tensor, weight_scale: torch.Tensor, logical_widths: List[int] +) -> Tuple[torch.Tensor, torch.Tensor]: + # Max scale to be used for requanitzation. + max_w_scale = weight_scale.max() + + # QKV / MLP is fused in the on disk checkpoint if any of the + # weight scales are still set to the default since we initialize + # N weight scales for N shards but we only load 1 weight scale + # from disk in this case. Skip requantization in this case (since) + # we already are quantized with the single scale. + # * Sample Model: nm-testing/Phi-3-mini-128k-instruct-FP8 + unfused_module_in_checkpoint = ( + weight_scale[-1] > torch.finfo(torch.float8_e4m3fn).min + ) + + # If unfused checkpoint, need requanize with the single scale. + if unfused_module_in_checkpoint: + start = 0 + for idx, logical_width in enumerate(logical_widths): + end = start + logical_width + weight_dq = per_tensor_dequantize(weight[start:end, :], weight_scale[idx]) + weight[start:end, :], _ = scaled_fp8_quant(weight_dq, max_w_scale) + start = end + + return max_w_scale, weight + + +def update_tensor_inplace(old: torch.Tensor, new: torch.Tensor) -> None: + old.copy_(new) + + +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/layer_utils.py +# Newly generated tensors need to replace existing tensors that are +# already registered as parameters by vLLM (and won't be freed) +def replace_parameter( + mod: torch.nn.Module, name: str, new: Union[torch.Tensor, torch.nn.Parameter] +) -> None: + + old = getattr(mod, name) + if ( + type(old) is type(new) + and old.dtype == new.dtype + and old.untyped_storage().nbytes() == new.untyped_storage().nbytes() + ): + # If we can just update in-place to avoid re-registering + # can be faster if the underlying storage is the same + update_tensor_inplace(old, new) + else: + # Fallback re-register parameter, convert to Parameter if necessary + # this not only ensures we don't register a tensor as a parameter, but + # also ensures that all parameter subclasses get re-registered as + # parameters for `torch.compile` compatibility + if not isinstance(new, torch.nn.Parameter): + new = torch.nn.Parameter(new, requires_grad=False) + mod.register_parameter(name, torch.nn.Parameter(new, requires_grad=False)) + + +def assert_fp8_all_close(a: torch.Tensor, b: torch.Tensor): + assert a.shape == b.shape + assert a.dtype == b.dtype == torch.float8_e4m3fn + + a_u8 = a.view(torch.uint8) + b_u8 = b.view(torch.uint8) + diff_u8 = (a_u8.to(torch.int16) - b_u8.to(torch.int16)).abs() + + numel = a.numel() + + count_diff_sign = ((a_u8 >= 0) & (b_u8 < 0)).sum().item() + count_tiny_diff = (diff_u8 >= 1).sum().item() + count_large_diff = (diff_u8 >= 2).sum().item() + + assert ( + (count_diff_sign == 0) + and (count_tiny_diff / numel < 0.005) + and (count_large_diff == 0) + ), f"{count_diff_sign=} {count_tiny_diff=} {count_large_diff=} {numel=}" + + +# Match dynamic rules with module name (prefix) and override quantize +# config if module (prefix) matches a rule +def override_config(config: QuantizationConfig, prefix: str): + weight_bits = get_dynamic_override(config, prefix, "bits", config.weight_bits) + if isinstance(weight_bits, int): + config.weight_bits = weight_bits + group_size = get_dynamic_override(config, prefix, "group_size", config.group_size) + if isinstance(group_size, int): + config.group_size = group_size + desc_act = get_dynamic_override(config, prefix, "desc_act", config.desc_act) + if isinstance(desc_act, bool): + config.desc_act = desc_act + + config.pack_factor = 32 // config.weight_bits # packed into int32 + if config.get_name() == "gptq_marlin": + is_sym = get_dynamic_override(config, prefix, "sym", config.is_sym) + if isinstance(is_sym, bool): + config.is_sym = is_sym + + if (config.weight_bits, config.is_sym) not in config.TYPE_MAP: + raise ValueError( + "Unsupported quantization config: " + f"bits={config.weight_bits}, sym={config.is_sym}" + ) + + config.quant_type = config.TYPE_MAP[(config.weight_bits, config.is_sym)] + elif config.get_name() == "gptq": + if config.weight_bits not in [2, 3, 4, 8]: + raise ValueError( + "Currently, only 2/3/4/8-bit weight quantization is " + f"supported for GPTQ, but got {config.weight_bits} bits." + ) + + +def get_dynamic_override( + config: QuantizationConfig, + layer_name: str, + key: Optional[str] = None, + default_value: Union[int, bool, None] = None, +) -> Union[Dict, int, bool, None]: + for pattern, pattern_dict in config.dynamic.items(): + # Negative match: matched modules are excluded from quantized init + if pattern.startswith("-:"): + if re.match(pattern.removeprefix("-:"), layer_name): + return False + # Positive match: matched modules have quant properties overrides + # base quant config + elif re.match(pattern.removeprefix("+:"), layer_name): + if key is None: + return pattern_dict + else: + return pattern_dict.get(key, default_value) + return default_value + + +def get_linear_quant_method( + config: QuantizationConfig, + layer: torch.nn.Module, + prefix: str, + linear_method_cls: type, +): + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.quantization.unquant import ( + UnquantizedEmbeddingMethod, + UnquantizedLinearMethod, + ) + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + + cloned_config = deepcopy(config) + parallel_lm_head_quantized = ( + isinstance(layer, ParallelLMHead) and cloned_config.lm_head_quantized + ) + + if isinstance(layer, LinearBase) or parallel_lm_head_quantized: + # False = skip module, None = no override, else = Positive match + if get_dynamic_override(cloned_config, layer_name=prefix) is False: + if parallel_lm_head_quantized: + return UnquantizedEmbeddingMethod() + return UnquantizedLinearMethod() + + if prefix: + # Dynamic per module/layer rules may override base config + override_config(cloned_config, prefix=prefix) + + return linear_method_cls(cloned_config) + return None + + +def get_pack_factor(num_bits): + assert 32 % num_bits == 0, f"Unsupported num_bits = {num_bits}" + return 32 // num_bits + + +def permute_rows( + q_w: torch.Tensor, + w_ref: torch.Tensor, + group_size: int, + test_perm: Optional[torch.Tensor] = None, +): + assert q_w.shape == w_ref.shape + + orig_device = q_w.device + k_size, _ = q_w.shape + + g_idx = torch.zeros((k_size,), dtype=torch.int32) + for i in range(k_size): + g_idx[i] = i // group_size + + # Simulate act_order by doing a random permutation on K + rand_perm = test_perm if test_perm is not None else torch.randperm(k_size) + + g_idx = g_idx[rand_perm].contiguous() + q_w = q_w[rand_perm, :].contiguous() + w_ref = w_ref[rand_perm, :].contiguous() + + return ( + w_ref.to(device=orig_device), + q_w.to(device=orig_device), + g_idx.to(device=orig_device), + rand_perm.to(device=orig_device), + ) + + +def pack_cols( + q_w: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + assert q_w.shape == (size_k, size_n) + + pack_factor = get_pack_factor(num_bits) + assert size_n % pack_factor == 0 + + orig_device = q_w.device + + q_w = q_w.cpu().numpy().astype(numpy.uint32) + + q_res = numpy.zeros((size_k, size_n // pack_factor), dtype=numpy.uint32) + + for i in range(pack_factor): + q_res |= q_w[:, i::pack_factor] << num_bits * i + + q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device) + q_res = q_res.contiguous() + + return q_res + + +def pack_rows( + q_w: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + assert q_w.shape == (size_k, size_n) + + pack_factor = get_pack_factor(num_bits) + assert size_k % pack_factor == 0 + + orig_device = q_w.device + + q_w = q_w.cpu().numpy().astype(numpy.uint32) + + q_res = numpy.zeros((size_k // pack_factor, size_n), dtype=numpy.uint32) + + for i in range(pack_factor): + q_res |= q_w[i::pack_factor, :] << num_bits * i + + q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device) + return q_res + + +def unpack_cols( + packed_q_w: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + pack_factor = get_pack_factor(num_bits) + assert size_n % pack_factor == 0 + assert packed_q_w.shape == ( + size_k, + size_n // pack_factor, + ), "packed_q_w.shape = {} size_k = {}, size_n = {} pack_Factor = {}".format( + packed_q_w.shape, size_k, size_n, pack_factor + ) + + orig_device = packed_q_w.device + + packed_q_w_cpu = packed_q_w.cpu().numpy().astype(numpy.uint32) + q_res = numpy.zeros((size_k, size_n), dtype=numpy.uint32) + + mask = (1 << num_bits) - 1 + for i in range(pack_factor): + vals = packed_q_w_cpu & mask + packed_q_w_cpu >>= num_bits + q_res[:, i::pack_factor] = vals + + q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device) + q_res = q_res.contiguous() + + return q_res + + +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/quant_utils.py +def quantize_weights( + w: torch.Tensor, + quant_type: ScalarType, + group_size: Optional[int], + zero_points: bool = False, + ref_zero_points_after_scales: bool = False, +): + assert ( + quant_type.is_integer() + ), "Floating point quantization may work but has not been tested" + assert not zero_points or group_size is not None, ( + "to have group zero points, group_size must be provided " + "(-1 group_size is channelwise)" + ) + + orig_device = w.device + orig_type = w.dtype + size_k, size_n = w.shape + + assert w.is_floating_point(), "w must be float" + + if group_size == -1: + group_size = size_k + + # Reshape to [groupsize, -1] + if group_size is not None and group_size < size_k: + w = w.reshape((-1, group_size, size_n)) + w = w.permute(1, 0, 2) + w = w.reshape((group_size, -1)) + + # Compute scale for each group + max_val = torch.max(w, 0, keepdim=True).values + min_val = torch.min(w, 0, keepdim=True).values + + max_q_val = quant_type.max() + min_q_val = quant_type.min() + + w_s = torch.Tensor([1.0]).to(w.device) # unscaled case + maybe_w_zp = None + if group_size is not None: + if zero_points: + assert not quant_type.is_signed() and quant_type.max() > 0 + w_s = (max_val - min_val).clamp(min=1e-5) / quant_type.max() + maybe_w_zp = ( + torch.round(torch.abs(min_val / w_s)).clamp(min_q_val, max_q_val).int() + ) + else: + # If the bias is such that there are no possible negative/positive + # values, set the max value to inf to avoid divide by 0 + w_s = torch.max( + abs(max_val / (max_q_val if max_q_val != 0 else torch.inf)), + abs(min_val / (min_q_val if min_q_val != 0 else torch.inf)), + ) + + # Quantize + w_q = torch.round(w / w_s).int() + (maybe_w_zp if zero_points else 0) + w_q = torch.clamp(w_q, min_q_val, max_q_val) + + # Compute ref (dequantized) + # For some kernels (namely Machete) the zero-points are applied after the + # scales are applied, for this case computing the reference in similar way + # allows us to use tighter error tolerances in our unit tests. + if ref_zero_points_after_scales and maybe_w_zp is not None: + w_ref = w_q.to(orig_type) * w_s - maybe_w_zp.to(orig_type) * w_s + else: + w_ref = (w_q - (maybe_w_zp if zero_points else 0)).to(orig_type) * w_s + + if quant_type.has_bias(): + w_q += quant_type.bias + + # Restore original shapes + if group_size is not None and group_size < size_k: + + def reshape_w(w): + w = w.reshape((group_size, -1, size_n)) + w = w.permute(1, 0, 2) + w = w.reshape((size_k, size_n)).contiguous() + return w + + w_q = reshape_w(w_q) + w_ref = reshape_w(w_ref) + w_s = w_s.reshape((-1, size_n)).contiguous() + + if maybe_w_zp is not None: + maybe_w_zp = maybe_w_zp.reshape((-1, size_n)).contiguous() + maybe_w_zp = maybe_w_zp.to(device=orig_device) + + return ( + w_ref.to(device=orig_device), + w_q.to(device=orig_device), + w_s if group_size is not None else None, + maybe_w_zp, + ) + + +SUPPORTED_GPTQ_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128] +SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] + + +def gptq_quantize_weights( + w: torch.Tensor, + quant_type: ScalarType, + group_size: int, + act_order: bool, + test_perm: Optional[torch.Tensor] = None, +): + size_k, _ = w.shape + + assert w.is_floating_point(), "w must be float" + assert ( + quant_type in SUPPORTED_GPTQ_QUANT_TYPES + ), f"Unsupported gptq type = {quant_type}" + assert group_size in SUPPORTED_GROUP_SIZES + [ + size_k + ], f"Unsupported groupsize = {group_size}" + + w_ref, w_q, w_s, _ = quantize_weights(w, quant_type, group_size) + + # Apply act_order + g_idx = torch.empty(0, dtype=torch.int, device=w.device) + rand_perm = torch.empty(0, dtype=torch.int, device=w.device) + if act_order: + assert ( + group_size < size_k + ), "For act_order, groupsize = {} must be less than size_k = {}".format( + group_size, size_k + ) + + w_ref, w_q, g_idx, rand_perm = permute_rows(w_q, w_ref, group_size, test_perm) + + return w_ref, w_q, w_s, g_idx, rand_perm + + +def sort_weights(q_w: torch.Tensor, g_idx: torch.Tensor): + orig_device = q_w.device + + sort_indices = torch.argsort(g_idx).to(dtype=torch.int32) # Sort based on g_idx + + g_idx = g_idx[sort_indices].contiguous() + q_w = q_w[sort_indices, :].contiguous() + + return ( + q_w.to(device=orig_device), + g_idx.to(device=orig_device), + sort_indices.to(device=orig_device), + ) diff --git a/python/sglang/srt/layers/quantization/w4afp8.py b/python/sglang/srt/layers/quantization/w4afp8.py new file mode 100644 index 000000000..e95247041 --- /dev/null +++ b/python/sglang/srt/layers/quantization/w4afp8.py @@ -0,0 +1,346 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional + +import torch +from torch.nn import Module +from torch.nn.parameter import Parameter + +from sglang.srt.distributed.parallel_state import get_moe_expert_parallel_world_size +from sglang.srt.layers.linear import LinearBase, UnquantizedLinearMethod +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.layers.quantization.utils import is_layer_skipped +from sglang.srt.utils import is_npu, set_weight_attrs + +_is_npu = is_npu() +if not _is_npu: + from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe + +if TYPE_CHECKING: + from sglang.srt.layers.moe import MoeRunnerConfig + from sglang.srt.layers.moe.ep_moe.layer import EPMoE + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +ACTIVATION_SCHEMES = ["static", "dynamic"] + +logger = logging.getLogger(__name__) + + +class W4AFp8Config(QuantizationConfig): + """Config class for MIXED_PRECISION W4AFp8.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = True, + is_checkpoint_w4afp8_serialized: bool = True, + linear_activation_scheme: str = "dynamic", + moe_activation_scheme: str = "static", + ignored_layers: Optional[List[str]] = None, + weight_block_size: Optional[List[int]] = None, + group_size: int = 128, + ) -> None: + super().__init__() + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + self.is_checkpoint_w4afp8_serialized = is_checkpoint_w4afp8_serialized + if is_checkpoint_w4afp8_serialized: + logger.warning("Detected w4afp8 checkpoint. Please note that") + if moe_activation_scheme not in ACTIVATION_SCHEMES: + raise ValueError(f"Unsupported activation scheme {moe_activation_scheme}") + self.linear_activation_scheme = linear_activation_scheme + self.moe_activation_scheme = moe_activation_scheme + self.ignored_layers = ignored_layers or [] + self.weight_block_size = [128, 128] + self.group_size = group_size + + @classmethod + def get_name(cls) -> str: + return "w4afp8" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.float8_e4m3fn] + + @classmethod + def get_min_capability(cls) -> int: + return 90 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> W4AFp8Config: + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_fp8_serialized = "fp8" in quant_method + is_checkpoint_w4afp8_serialized = "w4afp8" in quant_method + linear_activation_scheme = "dynamic" + moe_activation_scheme = "static" + weight_block_size = [128, 128] + return cls( + is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, + is_checkpoint_w4afp8_serialized=is_checkpoint_w4afp8_serialized, + linear_activation_scheme=linear_activation_scheme, + moe_activation_scheme=moe_activation_scheme, + weight_block_size=weight_block_size, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.ep_moe.layer import EPMoE + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.managers.schedule_batch import global_server_args_dict + + if isinstance(layer, LinearBase): + if is_layer_skipped(prefix, self.ignored_layers): + return UnquantizedLinearMethod() + return Fp8LinearMethod(self) + elif isinstance(layer, FusedMoE): + return W4AFp8MoEMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +def interleave_scales(scales: torch.Tensor) -> torch.Tensor: + """Interleave scales in groups of 4 similar to TRT-LLM implementation.""" + s_shape = scales.shape + # Reshape to separate groups of 4 + alignment = 4 if s_shape[2] % 4 == 0 else 1 + scales_interleaved = scales.reshape( + s_shape[0], s_shape[1], (s_shape[2] // alignment), alignment + ) + # Permute dimensions to interleave + scales_interleaved = scales_interleaved.permute(0, 2, 1, 3) + # Reshape back to original dimensions but with interleaved values + scales_interleaved = scales_interleaved.reshape( + s_shape[0], s_shape[2] // alignment, s_shape[1] * alignment + ) + return scales_interleaved.contiguous() + + +class W4AFp8MoEMethod(FusedMoEMethodBase): + def __init__(self, quant_config: W4AFp8Config): + self.quant_config = quant_config + + def create_weights( + self, + layer: EPMoE, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + assert "weight_loader" in extra_weight_attrs + + # Fused gate_up_proj (column parallel) + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition * 2, + hidden_size // 2, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + # down_proj (row parallel) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // 2, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} + ) + w13_weight_scale = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.quant_config.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale_inv", w13_weight_scale) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + w2_weight_scale = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size, + intermediate_size_per_partition // self.quant_config.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale_inv", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # Input scales + w13_input_scale = torch.nn.Parameter( + torch.ones((num_experts, 2), dtype=torch.bfloat16), + requires_grad=False, + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.bfloat16), + requires_grad=False, + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + + # Pre-populate the strides + device = layer.w13_weight.device + + self.a_strides1 = torch.full( + (num_experts, 3), + hidden_size, + device=device, + dtype=torch.int64, + ) + self.c_strides1 = torch.full( + (num_experts, 3), + 2 * intermediate_size_per_partition, + device=device, + dtype=torch.int64, + ) + self.a_strides2 = torch.full( + (num_experts, 3), + intermediate_size_per_partition, + device=device, + dtype=torch.int64, + ) + self.c_strides2 = torch.full( + (num_experts, 3), + hidden_size, + device=device, + dtype=torch.int64, + ) + self.b_strides1 = self.a_strides1 + self.s_strides13 = self.c_strides1 + self.b_strides2 = self.a_strides2 + self.s_strides2 = self.c_strides2 + + self.expert_offsets = torch.empty( + (num_experts + 1), dtype=torch.int32, device=device + ) + self.problem_sizes1 = torch.empty( + (num_experts, 3), dtype=torch.int32, device=device + ) + self.problem_sizes2 = torch.empty( + (num_experts, 3), dtype=torch.int32, device=device + ) + + return + + def process_weights_after_loading(self, layer: Module) -> None: + dtype = torch.bfloat16 + device = layer.w2_weight.device + + # Interleave w13_weight_scale (gate_up_proj) + w13_weight_scale = layer.w13_weight_scale_inv.to(dtype) + w13_weight_scale = interleave_scales(w13_weight_scale) + layer.w13_weight_scale_inv = Parameter(w13_weight_scale, requires_grad=False) + + # Interleave w2_weight_scale (down_proj) + w2_weight_scale = layer.w2_weight_scale_inv.to(dtype) + w2_weight_scale = interleave_scales(w2_weight_scale) + layer.w2_weight_scale_inv = Parameter(w2_weight_scale, requires_grad=False) + + # Process input scales + w13_input_scale_max = layer.w13_input_scale.max().to(dtype).item() + new_w13_input_scale = torch.tensor( + [w13_input_scale_max], + dtype=dtype, + device=device, + ) + layer.w13_input_scale = Parameter(new_w13_input_scale, requires_grad=False) + + w2_input_scale_max = layer.w2_input_scale.max().to(dtype).item() + new_w2_input_scale = torch.tensor( + [w2_input_scale_max], dtype=dtype, device=device + ) + layer.w2_input_scale = Parameter(new_w2_input_scale, requires_grad=False) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer: EPMoE, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + topk_weights, topk_ids, _ = topk_output + local_topk_ids = topk_ids + if get_moe_expert_parallel_world_size() > 1: + local_topk_ids = torch.where( + topk_ids == -1, + layer.num_experts, + topk_ids, + ) + + output = cutlass_w4a8_moe( + layer.start_expert_id, + layer.end_expert_id, + layer.num_experts, + x, + layer.w13_weight, + layer.w2_weight, + layer.w13_weight_scale_inv, + layer.w2_weight_scale_inv, + topk_weights, + topk_ids, + local_topk_ids, + self.a_strides1, + self.b_strides1, + self.c_strides1, + self.a_strides2, + self.b_strides2, + self.c_strides2, + self.s_strides13, + self.s_strides2, + self.expert_offsets, + self.problem_sizes1, + self.problem_sizes2, + layer.w13_input_scale, + layer.w2_input_scale, + ) + if self.moe_runner_config.routed_scaling_factor is not None: + output *= self.moe_runner_config.routed_scaling_factor + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/w8a8_fp8.py b/python/sglang/srt/layers/quantization/w8a8_fp8.py new file mode 100644 index 000000000..c68591fc6 --- /dev/null +++ b/python/sglang/srt/layers/quantization/w8a8_fp8.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import torch +from torch.nn.parameter import Parameter + +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.parameter import ChannelQuantScaleParameter, ModelWeightParameter +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8_kernel import ( + fp8_dtype, + is_fp8_fnuz, + per_token_group_quant_fp8, +) +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + cutlass_fp8_supported, + input_to_float8, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.utils import set_weight_attrs + +if TYPE_CHECKING: + from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +_is_fp8_fnuz = is_fp8_fnuz() + + +class W8A8Fp8Config(QuantizationConfig): + """Config class for W8A8 FP8 Quantization. + + Weight Quantization: + - Method: Static quantization + - Granularity: Per-channel + - Type: Symmetric + + Activation Quantization: + - Method: Dynamic quantization + - Granularity: Per-token + - Type: Symmetric + + Note: + - For models without offline quantization, weights will be quantized during model loading + - If CUTLASS is supported: Per-channel weight quantization is used + - If CUTLASS is not supported: Falls back to per-tensor weight quantization + """ + + def __init__(self, is_checkpoint_fp8_serialized: bool = False): + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.float16, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 89 + + @classmethod + def get_name(self) -> str: + return "w8a8_fp8" + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> W8A8Fp8Config: + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_fp8_serialized = ( + "compressed-tensors" in quant_method or "w8a8_fp8" in quant_method + ) + return cls(is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + return W8A8Fp8LinearMethod(self) + elif isinstance(layer, FusedMoE): + return W8A8FP8MoEMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class W8A8Fp8LinearMethod(LinearMethodBase): + + def __init__(self, quantization_config: W8A8Fp8Config): + self.cutlass_fp8_supported = cutlass_fp8_supported() + self.quantization_config = quantization_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight = layer.weight + + if self.quantization_config.is_checkpoint_fp8_serialized: + weight_scale = layer.weight_scale.detach() + # If checkpoint offline quantized with w8a8_fp8, load the weight and weight_scale directly. + if _is_fp8_fnuz: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, weight_scale=weight_scale + ) + + layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + else: + # If checkpoint not offline quantized, quantize the weights with per-channel quantization. + if self.cutlass_fp8_supported: + # if cutlass supported, we use cutlass_scaled_mm + # which requires per-channel quantization on weight + qweight, weight_scale = per_token_group_quant_fp8( + layer.weight, layer.weight.shape[-1] + ) + weight_scale = weight_scale.t().contiguous() + else: + # if cutlass not supported, we fall back to use torch._scaled_mm + # which requires per tensor quantization on weight + qweight, weight_scale = input_to_float8(layer.weight, dtype=fp8_dtype) + + # Update the layer with the new values. + layer.weight = Parameter(qweight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + layer.input_scale = None + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + weight_dtype = ( + torch.float8_e4m3fn + if self.quantization_config.is_checkpoint_fp8_serialized + else params_dtype + ) + + weight_loader = extra_weight_attrs.get("weight_loader") + self.logical_widths = output_partition_sizes + + weight = ModelWeightParameter( + data=torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + if self.quantization_config.is_checkpoint_fp8_serialized: + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + else: + layer.weight_scale = None + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + return apply_fp8_linear( + x, + layer.weight, + layer.weight_scale, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + ) + + +class W8A8FP8MoEMethod(FusedMoEMethodBase): + """MoE method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: W8A8Fp8Config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=fp8_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=fp8_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + w13_input_scale = None + layer.register_parameter("w13_input_scale", w13_input_scale) + + w2_input_scale = None + layer.register_parameter("w2_input_scale", w2_input_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.w13_weight = Parameter(layer.w13_weight, requires_grad=False) + layer.w2_weight = Parameter(layer.w2_weight, requires_grad=False) + layer.w13_weight_scale = Parameter( + layer.w13_weight_scale.data, requires_grad=False + ) + layer.w2_weight_scale = Parameter( + layer.w2_weight_scale.data, requires_grad=False + ) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + use_fp8_w8a8=True, + per_channel_quant=True, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + ) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/w8a8_int8.py b/python/sglang/srt/layers/quantization/w8a8_int8.py new file mode 100644 index 000000000..5ccb0259d --- /dev/null +++ b/python/sglang/srt/layers/quantization/w8a8_int8.py @@ -0,0 +1,1034 @@ +from __future__ import annotations + +import importlib +import sys +from types import MappingProxyType +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Tuple, + Union, + cast, +) + +import torch +from torch.nn.parameter import Parameter + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.parameter import ( + ChannelQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer +from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8 +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.utils import ( + apply_module_patch, + cpu_has_amx_support, + is_cpu, + is_cuda, + is_npu, + set_weight_attrs, + use_intel_amx_backend, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +_is_cuda = is_cuda() +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() +if _is_cuda: + from sgl_kernel import int8_scaled_mm +_is_npu = is_npu() + +if _is_npu: + import torch_npu + + try: + from mindie_turbo import _ops as ops + from mindie_turbo.quantize.quant_utils import quant_per_tensor + except ImportError: + useMindIETurbo = False + else: + useMindIETurbo = True + + +# func refers to RMSNorm.__init__ +def npu_wrapper_rmsnorm_init(func): + def init(self, hidden_size: int, **extra_args) -> None: + func(self, hidden_size, **extra_args) + self.ignore_anti = True + # The Ascend w8a8_int8 quantization requires adding a bias in rmsnorm + self.bias = torch.nn.Parameter(torch.zeros(hidden_size), requires_grad=False) + + return init + + +# func refers to RMSNorm.forward_oot +def npu_wrapper_rmsnorm_forward(func): + def _rmsnorm_forward_oot( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if not x.is_contiguous(): + x = x.contiguous() + if residual is not None: + out, _, residual_out = torch_npu.npu_add_rms_norm( + residual, x, self.weight.data, self.variance_epsilon + ) + out = out + self.bias + return out.to(x.dtype), residual_out + + out = torch_npu.npu_rms_norm(x, self.weight.data, self.variance_epsilon)[0] + out = out + self.bias + return out.to(x.dtype) + + return _rmsnorm_forward_oot + + +def npu_fused_experts( + hidden_states: torch.Tensor, + w13: torch.Tensor, + w13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, +): + original_shape = hidden_states.shape + original_dtype = hidden_states.dtype + scale_dtype = original_dtype if original_dtype == torch.bfloat16 else torch.float32 + if len(original_shape) == 3: + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + num_tokens = hidden_states.shape[0] + num_experts = w13.shape[0] + row_idx_len = num_tokens * top_k + row_idx = ( + torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device) + .view(top_k, -1) + .permute(1, 0) + .contiguous() + ) + hidden_states, expanded_row_idx, expanded_expert_idx = ( + torch_npu.npu_moe_init_routing( + hidden_states, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens + ) + ) + expert_tokens = torch_npu.npu_moe_compute_expert_tokens( + expanded_expert_idx, num_experts + ) + expert_tokens = expert_tokens.to(torch.int64) + # gmm1: gate_up_proj + hidden_states, pertoken_scale = torch_npu.npu_dynamic_quant(hidden_states) + hidden_states = torch_npu.npu_grouped_matmul( + x=[hidden_states], + weight=[w13], + scale=[w13_scale.to(scale_dtype)], + per_token_scale=[pertoken_scale], + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + # act_fn: swiglu + hidden_states = torch_npu.npu_swiglu(hidden_states) + hidden_states, pertoken_scale = torch_npu.npu_dynamic_quant(hidden_states) + # gmm2: down_proj + hidden_states = torch_npu.npu_grouped_matmul( + x=[hidden_states], + weight=[w2], + scale=[w2_scale.to(scale_dtype)], + per_token_scale=[pertoken_scale], + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + + final_hidden_states = torch_npu.npu_moe_finalize_routing( + hidden_states, + skip1=None, + skip2=None, + bias=None, + scales=topk_weights, + expanded_src_to_dst_row=expanded_row_idx, + export_for_source_row=topk_ids, + ) + if len(original_shape) == 3: + final_hidden_states = final_hidden_states.view(original_shape) + return final_hidden_states + + +class W8A8Int8Config(QuantizationConfig): + """Config class for W8A8 Int8 Quantization. + + - Weight: static, per-channel, symmetric + - Activation: dynamic, per-token, symmetric + """ + + def __init__(self, quant_config: Dict[str, Any] = {}): + super().__init__() + self.quant_description = quant_config + self.is_dynamic = quant_config.get("is_dynamic", False) + ignore = cast(List[str], quant_config.get("ignore", [])) + self.ignore = ignore if ignore is not None else [] + packed_modules_mapping = quant_config.get("packed_modules_mapping", {}) + self.packed_modules_mapping = ( + packed_modules_mapping if packed_modules_mapping is not None else {} + ) + + if _is_npu: + # Ascend w8a8_int8 quantization with bias, use wrappers to isolate the effects between models + for name in self.quant_description.keys(): + if "norm.bias" in name: + apply_module_patch( + "sglang.srt.layers.layernorm.RMSNorm", + "__init__", + [npu_wrapper_rmsnorm_init], + ) + apply_module_patch( + "sglang.srt.layers.layernorm.RMSNorm", + "forward_npu", + [npu_wrapper_rmsnorm_forward], + ) + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return ( + [torch.float16, torch.bfloat16] + if not _is_npu + else [torch.int8, torch.float16, torch.bfloat16] + ) + + @classmethod + def get_min_capability(cls) -> int: + if _is_npu: + raise NotImplementedError( + 'NPU hardware does not support "get_min_capability" feature.' + ) + else: + return 75 + + @classmethod + def get_name(self) -> str: + return "w8a8_int8" + + @classmethod + def get_config_filenames(cls) -> List[str]: + filenames = [] + if _is_npu: + filenames.append("quant_model_description.json") + return filenames + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> W8A8Int8Config: + return cls(config) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if _is_npu: + if isinstance(layer, LinearBase): + key = "model" + if "vision_model" in prefix: + key = "vision_model" + elif "visual" in prefix: + key = "visual" + packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {}) + prefix_in_quant_config = prefix + proj_name = prefix.split(".")[-1] + if proj_name in packed_modules_mapping_subset: + prefix_in_quant_config = prefix.replace( + proj_name, packed_modules_mapping_subset[proj_name][0] + ) + self.is_dynamic = ( + self.quant_description[prefix_in_quant_config + ".weight"] + == "W8A8_DYNAMIC" + ) + if self.is_layer_skipped(prefix, packed_modules_mapping_subset): + return UnquantizedLinearMethod() + return ( + NPU_W8A8DynamicLinearMethod(self) + if self.is_dynamic + else NPU_W8A8LinearMethod(self) + ) + elif isinstance(layer, FusedMoE): + return NPU_W8A8MoEMethod(self) + return None + + if should_ignore_layer( + prefix, ignore=self.ignore, fused_mapping=self.packed_modules_mapping + ): + return UnquantizedLinearMethod() + if isinstance(layer, LinearBase): + return W8A8Int8LinearMethod(self) + elif isinstance(layer, FusedMoE): + return W8A8Int8MoEMethod(self) + return None + + def is_layer_skipped( + self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({}) + ): + # adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped + proj_name = prefix.split(".")[-1] + if proj_name in fused_mapping: + shard_prefixes = [ + prefix.replace(proj_name, shard_proj_name) + for shard_proj_name in fused_mapping[proj_name] + ] + + is_skipped = None + for shard_prefix in shard_prefixes: + is_shard_skipped = ( + self.quant_description[shard_prefix + ".weight"] == "FLOAT" + ) + + if is_skipped is None: + is_skipped = is_shard_skipped + elif is_shard_skipped != is_skipped: + raise ValueError( + f"Detected some but not all shards of {prefix} " + "are quantized. All shards of fused layers " + "to have the same precision." + ) + else: + is_skipped = self.quant_description[prefix + ".weight"] == "FLOAT" + + assert is_skipped is not None + return is_skipped + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class W8A8Int8LinearMethod(LinearMethodBase): + + def __init__(self, quantization_config: W8A8Int8Config): + self.quantization_config = quantization_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if _is_cpu: + assert ( + _is_cpu_amx_available + ), "W8A8Int8LinearMethod on CPU requires that CPU has AMX support" + _amx_process_weight_after_loading(layer, ["weight"]) + else: + layer.weight = Parameter(layer.weight.t(), requires_grad=False) + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + + weight_loader = extra_weight_attrs.get("weight_loader") + self.logical_widths = output_partition_sizes + + weight = ModelWeightParameter( + data=torch.empty( + sum(output_partition_sizes), input_size_per_partition, dtype=torch.int8 + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + if use_intel_amx_backend(layer): + return torch.ops.sgl_kernel.int8_scaled_mm_with_quant( + x, + layer.weight, + layer.weight_scale, + bias, + x.dtype, + True, # is_vnni + ) + + x_q, x_scale = per_token_quant_int8(x) + + return int8_scaled_mm( + x_q, layer.weight, x_scale, layer.weight_scale, out_dtype=x.dtype, bias=bias + ) + + +class W8A8Int8MoEMethod(FusedMoEMethodBase): + """MoE method for INT8. + Supports loading INT8 checkpoints with static weight scale and + dynamic/static activation scale. + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: W8A8Int8Config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + tp_size = get_tensor_model_parallel_world_size() + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + w13_input_scale = None + layer.register_parameter("w13_input_scale", w13_input_scale) + + w2_input_scale = None + layer.register_parameter("w2_input_scale", w2_input_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if _is_cpu: + assert ( + _is_cpu_amx_available + ), "W8A8Int8MoEMethod on CPU requires that CPU has AMX support" + _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"]) + else: + layer.w13_weight = Parameter(layer.w13_weight, requires_grad=False) + layer.w2_weight = Parameter(layer.w2_weight, requires_grad=False) + layer.w13_weight_scale = Parameter( + layer.w13_weight_scale.data, requires_grad=False + ) + layer.w2_weight_scale = Parameter( + layer.w2_weight_scale.data, requires_grad=False + ) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> torch.Tensor: + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + if use_intel_amx_backend(layer): + from sglang.srt.layers.moe.topk import apply_topk_weights_cpu + + topk_weights, topk_ids, _ = topk_output + x, topk_weights = apply_topk_weights_cpu( + self.moe_runner_config.apply_router_weight_on_input, topk_weights, x + ) + output = torch.ops.sgl_kernel.fused_experts_cpu( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + False, # inplace See [Note] inplace should be False in fused_experts. + True, # use_int8_w8a8 + False, # use_fp8_w8a16 + layer.w13_weight_scale, # w1_scale + layer.w2_weight_scale, # w2_scale + None, # block_size + layer.w13_input_scale, # a1_scale + layer.w2_input_scale, # a2_scale + True, # is_vnni + ) + return StandardCombineInput(hidden_states=output) + + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + use_int8_w8a8=True, + per_channel_quant=True, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + ) + return self.runner.run(dispatch_output, quant_info) + + +class NPU_W8A8LinearMethodImpl: + """Linear method for NPU W8A8.""" + + def __init__(self) -> None: + # aclnn quant matmul requires to transpose matrix B, set to true by default. + self.transpose_weight = True + + @staticmethod + def get_weight( + input_size: int, + output_size: int, + params_dtype: torch.dtype = torch.bfloat16, + ) -> Dict[str, Any]: + params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)} + return params_dict + + @staticmethod + def get_pertensor_param(params_dtype: torch.dtype) -> Dict[str, Any]: + params_dict = {} + params_dict["input_scale"] = torch.empty(1, dtype=params_dtype) + params_dict["input_offset"] = torch.empty(1, dtype=params_dtype) + return params_dict + + @staticmethod + def get_perchannel_param( + output_size: int, + params_dtype: torch.dtype, + ) -> Dict[str, Any]: + params_dict = {} + params_dict["quant_bias"] = torch.empty(output_size, dtype=torch.int32) + if params_dtype == torch.bfloat16: + params_dict["deq_scale"] = torch.empty(output_size, dtype=torch.float32) + elif params_dtype == torch.float16: + params_dict["deq_scale"] = torch.empty(output_size, dtype=torch.int64) + params_dict["weight_scale"] = torch.empty(output_size, 1, dtype=params_dtype) + params_dict["weight_offset"] = torch.empty(output_size, 1, dtype=params_dtype) + return params_dict + + @staticmethod + def apply( + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # To prevent import loops + from sglang.srt.layers.linear import RowParallelLinear + + original_dtype = x.dtype + if original_dtype != torch.int8: + x = torch_npu.npu_quantize( + x, + layer.aclnn_input_scale_reciprocal, + layer.aclnn_input_offset, + torch.qint8, + -1, + False, + ) + # Only fuse bias add into GEMM for rank 0 (this ensures that + # bias will not get added more than once in Attention TP>1 case) + if isinstance(layer, RowParallelLinear) and layer.tp_rank > 0: + quant_bias = None + else: + quant_bias = layer.quant_bias + return torch_npu.npu_quant_matmul( + x, + layer.weight, + layer.deq_scale, + bias=quant_bias, + output_dtype=original_dtype, + ) + + def process_weights_after_loading(self, layer): + expanding_factor = layer.weight.data.shape[1] + layer.aclnn_input_scale = torch.nn.Parameter( + layer.input_scale.data.repeat(expanding_factor).to(device="npu"), + requires_grad=False, + ) + layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter( + layer.input_scale.data.repeat(expanding_factor).to(device="npu"), + requires_grad=False, + ) + layer.aclnn_input_offset = torch.nn.Parameter( + layer.input_offset.data.repeat(expanding_factor).to(device="npu"), + requires_grad=False, + ) + if self.transpose_weight: + layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() + layer.weight_scale.data = torch.flatten(layer.weight_scale.data) + layer.weight_offset.data = torch.flatten(layer.weight_offset.data) + + +class NPU_W8A8LinearMethodMTImpl: + """Linear method for NPU W8A8.""" + + def __init__(self) -> None: + self.transpose_weight = True + + @staticmethod + def get_weight( + input_size: int, + output_size: int, + params_dtype: torch.dtype = torch.bfloat16, + ) -> Dict[str, Any]: + params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)} + return params_dict + + @staticmethod + def get_pertensor_param(params_dtype: torch.dtype) -> Dict[str, Any]: + params_dict = {} + params_dict["input_scale"] = torch.empty(1, dtype=params_dtype) + params_dict["input_offset"] = torch.empty(1, dtype=torch.int8) + return params_dict + + @staticmethod + def get_perchannel_param( + output_size: int, + params_dtype: torch.dtype, + ) -> Dict[str, Any]: + params_dict = {} + params_dict["quant_bias"] = torch.empty(output_size, dtype=torch.int32) + if params_dtype == torch.bfloat16: + params_dict["deq_scale"] = torch.empty(output_size, dtype=torch.float32) + elif params_dtype == torch.float16: + params_dict["deq_scale"] = torch.empty(output_size, dtype=torch.int64) + params_dict["weight_scale"] = torch.empty(output_size, 1, dtype=params_dtype) + params_dict["weight_offset"] = torch.empty(output_size, 1, dtype=params_dtype) + return params_dict + + @staticmethod + def apply( + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # To prevent import loops + from sglang.srt.layers.linear import RowParallelLinear + + original_dtype = x.dtype + if original_dtype != torch.int8: + x = quant_per_tensor(x, layer.input_scale, layer.input_offset) + + # Only fuse bias add into GEMM for rank 0 (this ensures that + # bias will not get added more than once in Attention TP>1 case) + if isinstance(layer, RowParallelLinear) and layer.tp_rank > 0: + quant_bias = None + else: + quant_bias = layer.quant_bias + + return ops.quant_matmul( + x=x, weight=layer.weight, deq_scale=layer.deq_scale, deq_bias=quant_bias + ) + + def process_weights_after_loading(self, layer): + layer.aclnn_deq_scale = torch.nn.Parameter( + torch_npu.npu_trans_quant_param(layer.deq_scale.npu()).to(device="npu"), + requires_grad=False, + ) + + +class NPU_W8A8LinearMethod(LinearMethodBase): + """Linear method for NPU quantization. + + This class search for specific quantization + implementation supported on NPU hardware for linear methods. + + Args: + quant_config: The NPU quantization config. + """ + + def __init__(self, quantization_config: W8A8Int8Config) -> None: + self.quantization_config = quantization_config + self.quant_method = ( + NPU_W8A8LinearMethodMTImpl() + if useMindIETurbo + else NPU_W8A8LinearMethodImpl() + ) + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + weight_dict = self.quant_method.get_weight( + input_size_per_partition, output_size_per_partition, params_dtype + ) + for weight_name, weight_param in weight_dict.items(): + param = torch.nn.Parameter(weight_param, requires_grad=False) + set_weight_attrs(param, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter(weight_name, param) + set_weight_attrs(param, extra_weight_attrs) + + pertensor_dict = self.quant_method.get_pertensor_param(params_dtype) + for pertensor_name, pertensor_param in pertensor_dict.items(): + param = PerTensorScaleParameter( + data=pertensor_param, weight_loader=weight_loader + ) + # disable warning + param.ignore_warning = True + layer.register_parameter(pertensor_name, param) + + perchannel_dict = self.quant_method.get_perchannel_param( + output_size_per_partition, params_dtype + ) + for perchannel_name, perchannel_param in perchannel_dict.items(): + param = torch.nn.Parameter(perchannel_param, requires_grad=False) + set_weight_attrs(param, {"output_dim": 0}) + layer.register_parameter(perchannel_name, param) + set_weight_attrs(param, extra_weight_attrs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if hasattr(self.quant_method, "process_weights_after_loading"): + self.quant_method.process_weights_after_loading(layer) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.quant_method.apply(layer, x, bias) + + +class NPU_W8A8DynamicLinearMethodImpl: + """Linear method for NPU W8A8_DYNAMIC.""" + + def __init__(self): + self.transpose_weight = True + + @staticmethod + def get_weight( + input_size: int, output_size: int, params_dtype: torch.dtype + ) -> Dict[str, Any]: + params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)} + return params_dict + + @staticmethod + def get_pertensor_param(params_dtype: torch.dtype) -> Dict[str, Any]: + return {} + + @staticmethod + def get_perchannel_param( + output_size: int, + params_dtype: torch.dtype, + ) -> Dict[str, Any]: + params_dict = {} + params_dict["weight_scale"] = torch.empty(output_size, 1, dtype=params_dtype) + params_dict["weight_offset"] = torch.empty(output_size, 1, dtype=params_dtype) + return params_dict + + @staticmethod + def apply( + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + tp_rank: Optional[int] = 0, + ) -> torch.Tensor: + original_dtype = x.dtype + quant_out, dynamic_scale = torch_npu.npu_dynamic_quant(x) + return torch_npu.npu_quant_matmul( + quant_out, + layer.weight, + layer.weight_scale, + pertoken_scale=dynamic_scale, + bias=bias, + output_dtype=original_dtype, + ) + + def process_weights_after_loading(self, layer): + if self.transpose_weight: + layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() + layer.weight_scale.data = layer.weight_scale.data.flatten() + layer.weight_scale_fp32 = layer.weight_scale.data.to(torch.float32) + layer.weight_offset.data = layer.weight_offset.data.flatten() + + +class NPU_W8A8DynamicLinearMethod(LinearMethodBase): + """Linear method for NPU quantization. + + This class search for specific quantization + implementations supported on NPU hardware for linear methods. + + Args: + quant_config: The NPU quantization config. + """ + + def __init__(self, quantization_config: W8A8Int8Config) -> None: + self.quantization_config = quantization_config + self.quant_method = NPU_W8A8DynamicLinearMethodImpl() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + weight_dict = self.quant_method.get_weight( + input_size_per_partition, output_size_per_partition, params_dtype + ) + for weight_name, weight_param in weight_dict.items(): + param = torch.nn.Parameter(weight_param, requires_grad=False) + set_weight_attrs(param, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter(weight_name, param) + set_weight_attrs(param, extra_weight_attrs) + + pertensor_dict = self.quant_method.get_pertensor_param(params_dtype) + for pertensor_name, pertensor_param in pertensor_dict.items(): + param = PerTensorScaleParameter( + data=pertensor_param, weight_loader=weight_loader + ) + # disable warning + param.ignore_warning = True + layer.register_parameter(pertensor_name, param) + + perchannel_dict = self.quant_method.get_perchannel_param( + output_size_per_partition, params_dtype + ) + for perchannel_name, perchannel_param in perchannel_dict.items(): + param = torch.nn.Parameter(perchannel_param, requires_grad=False) + set_weight_attrs(param, {"output_dim": 0}) + layer.register_parameter(perchannel_name, param) + set_weight_attrs(param, extra_weight_attrs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if hasattr(self.quant_method, "process_weights_after_loading"): + self.quant_method.process_weights_after_loading(layer) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.quant_method.apply(layer, x, bias) + + +class NPU_W8A8MoEMethod(FusedMoEMethodBase): + """MoE method for NPU quantization. + + This class search for specific quantization + implementations supported on NPU hardware for moe methods. + + Args: + quant_config: The NPU quantization config. + """ + + def __init__(self, quantization_config: W8A8Int8Config) -> None: + self.quantization_config = quantization_config + self.quant_method = self + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + self.num_experts = num_experts + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + + # weight + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + # scale + w13_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + w2_weight_scale = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + # offset + w13_weight_offset = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_offset", w13_weight_offset) + set_weight_attrs(w13_weight_offset, extra_weight_attrs) + w2_weight_offset = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_offset", w2_weight_offset) + set_weight_attrs(w2_weight_offset, extra_weight_attrs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.w13_weight = Parameter( + layer.w13_weight.data.transpose(1, 2).contiguous(), requires_grad=False + ) + layer.w2_weight = Parameter( + layer.w2_weight.data.transpose(1, 2).contiguous(), requires_grad=False + ) + layer.w13_weight_scale = Parameter( + layer.w13_weight_scale.data.squeeze(-1).contiguous(), requires_grad=False + ) + layer.w2_weight_scale = Parameter( + layer.w2_weight_scale.data.squeeze(-1).contiguous(), requires_grad=False + ) + layer.w13_weight_offset = Parameter( + layer.w13_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False + ) + layer.w2_weight_offset = Parameter( + layer.w2_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False + ) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + topk_weights, topk_ids, _ = topk_output + topk_ids = topk_ids.to(torch.int32) + topk_weights = topk_weights.to(x.dtype) + output = npu_fused_experts( + hidden_states=x, + w13=layer.w13_weight, + w13_scale=layer.w13_weight_scale, + w2=layer.w2_weight, + w2_scale=layer.w2_weight_scale, + topk_weights=topk_weights, + topk_ids=topk_ids, + top_k=topk_ids.shape[1], + ) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py new file mode 100644 index 000000000..0719cdd29 --- /dev/null +++ b/python/sglang/srt/layers/radix_attention.py @@ -0,0 +1,116 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""Radix attention.""" +from __future__ import annotations + +from enum import Enum +from typing import TYPE_CHECKING, Optional + +from torch import nn + +if TYPE_CHECKING: + from sglang.srt.layers.quantization.base_config import QuantizationConfig + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +class AttentionType(Enum): + """ + Attention type. + Use string to be compatible with `torch.compile`. + """ + + # Decoder attention between previous layer Q/K/V + DECODER = "decoder" + # Encoder attention between previous layer Q/K/V + ENCODER_ONLY = "encoder_only" + + +class RadixAttention(nn.Module): + """ + The attention layer implementation. + """ + + def __init__( + self, + num_heads: int, + head_dim: int, + scaling: float, + num_kv_heads: int, + layer_id: int, + logit_cap: float = 0.0, + v_head_dim: int = -1, + sliding_window_size: int = -1, + is_cross_attention: bool = False, + pos_encoding_mode: str = "NONE", + logit_capping_method: str = "tanh", + quant_config: Optional[QuantizationConfig] = None, + attn_type: AttentionType = AttentionType.DECODER, + use_irope: bool = False, + prefix: str = "", + ): + super().__init__() + self.tp_q_head_num = num_heads + self.tp_k_head_num = num_kv_heads + self.tp_v_head_num = num_kv_heads + self.head_dim = head_dim + self.qk_head_dim = head_dim + self.v_head_dim = v_head_dim if v_head_dim != -1 else head_dim + self.scaling = scaling + self.layer_id = layer_id + self.logit_cap = logit_cap + self.sliding_window_size = sliding_window_size or -1 + self.is_cross_attention = is_cross_attention + self.use_irope = use_irope + self.k_scale = None + self.v_scale = None + self.k_scale_float = None + self.v_scale_float = None + self.quant_method = None + if quant_config is not None: + self.quant_method = quant_config.get_quant_method(self, prefix=prefix) + if self.quant_method is not None: + self.quant_method.create_weights(self) + self.attn_type = attn_type + + self.pos_encoding_mode = pos_encoding_mode + self.logit_capping_method = logit_capping_method + self.xai_temperature_len = -1 + + def forward( + self, + q, + k, + v, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + if k is not None: + # For cross-layer sharing, kv can be None + assert v is not None + if "k_rope" not in kwargs: + k = k.view(-1, self.tp_k_head_num, self.qk_head_dim) + v = v.view(-1, self.tp_v_head_num, self.v_head_dim) + else: + k = k.view(-1, self.tp_k_head_num, self.v_head_dim) + + return forward_batch.attn_backend.forward( + q, + k, + v, + self, + forward_batch, + save_kv_cache, + **kwargs, + ) diff --git a/python/sglang/srt/layers/rocm_linear_utils.py b/python/sglang/srt/layers/rocm_linear_utils.py new file mode 100644 index 000000000..ee7dd1f59 --- /dev/null +++ b/python/sglang/srt/layers/rocm_linear_utils.py @@ -0,0 +1,44 @@ +import torch +from aiter.ops.triton.fused_qk_concat import fused_qk_rope_cat +from aiter.ops.triton.gemm_a16w16 import gemm_a16w16 +from aiter.ops.triton.gemm_a16w16_atomic import gemm_a16w16_atomic + +from sglang.srt.utils import BumpAllocator + +__all__ = ["fused_qk_rope_cat"] + + +def aiter_dsv3_router_gemm( + hidden_states: torch.Tensor, + weight: torch.Tensor, + gemm_output_zero_allocator: BumpAllocator = None, +): + M = hidden_states.shape[0] + N = weight.shape[0] + y = None + + if M <= 256: + # TODO (cagri): convert to bfloat16 as part of another kernel to save time + # for now it is also coupled with zero allocator. + if gemm_output_zero_allocator != None: + y = gemm_output_zero_allocator.allocate(M * N).view(M, N) + else: + y = torch.zeros((M, N), dtype=torch.float32, device=hidden_states.device) + + if y is not None: + logits = gemm_a16w16_atomic(hidden_states, weight, y=y).to(hidden_states.dtype) + else: + logits = gemm_a16w16(hidden_states, weight) + + return logits + + +def get_dsv3_gemm_output_zero_allocator_size( + n_routed_experts: int, num_moe_layers: int, allocate_size: int, embedding_dim: int +): + if embedding_dim != 7168 or n_routed_experts != 256: + return 0 + + per_layer_size = 256 * (allocate_size + n_routed_experts) + + return num_moe_layers * per_layer_size diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py new file mode 100644 index 000000000..05f068557 --- /dev/null +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -0,0 +1,2017 @@ +# Adapted from https://raw.githubusercontent.com/vllm-project/vllm/refs/tags/v0.6.6.post1/vllm/model_executor/layers/rotary_embedding.py + +"""Rotary Positional Embeddings.""" +import itertools +import math +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +from sglang.srt.custom_op import CustomOp +from sglang.srt.utils import ( + cpu_has_amx_support, + get_bool_env_var, + is_cpu, + is_cuda, + is_hip, + is_npu, +) + +_is_cuda = is_cuda() +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_is_npu = is_npu() +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() + +if _is_cuda: + from sgl_kernel import apply_rope_with_cos_sin_cache_inplace +if _use_aiter: + from aiter.rotary_embedding import get_rope as aiter_get_rope + +if is_npu(): + import torch_npu + + +def _rotate_neox(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _rotate_gptj(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., ::2] + x2 = x[..., 1::2] + x = torch.stack((-x2, x1), dim=-1) + return x.flatten(-2) + + +def _apply_rotary_emb( + 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) + + +class RotaryEmbedding(CustomOp): + """Original rotary positional embedding.""" + + 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() + # NOTE(ByronHsu): cache needs to be in FP32 for numerical stability + if not _is_cuda: + cache = cache.to(dtype) + + if ( + not (_is_cuda or _is_npu) or self.head_size not in [64, 128, 256, 512] + ) and not (_is_cpu and _is_cpu_amx_available): + from vllm._custom_ops import rotary_embedding + + self.vllm_rotary_embedding = rotary_embedding + + 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: + """Compute the inverse frequency.""" + # NOTE(woosuk): To exactly match the HF implementation, we need to + # use CPU to compute the cache and then move it to GPU. However, we + # create the cache on GPU for faster initialization. This may cause + # a slight numerical difference between the HF implementation and ours. + 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 forward_native( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """A PyTorch-native implementation of forward().""" + if offsets is not None: + positions = positions + offsets + positions = positions.flatten() + num_tokens = positions.shape[0] + cos_sin = self.cos_sin_cache.index_select(0, positions) + cos, sin = cos_sin.chunk(2, dim=-1) + + query_shape = query.shape + query = query.view(num_tokens, -1, self.head_size) + query_rot = query[..., : self.rotary_dim] + query_pass = query[..., self.rotary_dim :] + query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style) + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) + + key_shape = key.shape + key = key.view(num_tokens, -1, self.head_size) + key_rot = key[..., : self.rotary_dim] + key_pass = key[..., self.rotary_dim :] + key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style) + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) + return query, key + + def forward_npu( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """A PyTorch-npu implementation of forward().""" + import os + + if get_bool_env_var("SGLANG_ENABLE_TORCH_COMPILE"): + return self.forward_native(positions, query, key, offsets) + else: + rotary_mode = "half" + if self.is_neox_style: + rotary_mode = "half" + else: + rotary_mode = "interleave" + mrope_section = [0, 0, 0] + query_out, key_out = torch_npu.npu_mrope( + positions, + query, + key, + self.cos_sin_cache, + self.head_size, + mrope_section=mrope_section, + rotary_mode=rotary_mode, + ) + return query_out, key_out + + def forward_cpu( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + positions = torch.add(positions, offsets) if offsets is not None else positions + if _is_cpu_amx_available: + return torch.ops.sgl_kernel.rotary_embedding_cpu( + positions, + query, + key, + self.head_size, + self.cos_sin_cache, + self.is_neox_style, + ) + else: + return self.forward_native(positions, query, key, offsets) + + def forward_cuda( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + fused_set_kv_buffer_arg=None, # Optional[FusedSetKVBufferArg] + ) -> Tuple[torch.Tensor, torch.Tensor]: + if _is_cuda and (self.head_size in [64, 128, 256, 512]): + apply_rope_with_cos_sin_cache_inplace( + positions=positions, + query=query, + key=key, + head_size=self.head_size, + cos_sin_cache=self.cos_sin_cache, + is_neox=self.is_neox_style, + # Compatible with old sgl-kernel + **( + dict(fused_set_kv_buffer_arg=fused_set_kv_buffer_arg) + if fused_set_kv_buffer_arg is not None + else {} + ), + ) + else: + assert ( + fused_set_kv_buffer_arg is None + ), "save kv cache is not supported for vllm_rotary_embedding." + self.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype) + self.vllm_rotary_embedding( + positions, + query, + key, + self.head_size, + self.cos_sin_cache, + self.is_neox_style, + ) + return query, key + + def extra_repr(self) -> str: + s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" + s += f", max_position_embeddings={self.max_position_embeddings}" + s += f", base={self.base}, is_neox_style={self.is_neox_style}" + return s + + +class LinearScalingRotaryEmbedding(RotaryEmbedding): + """RotaryEmbedding extended with linear scaling. + + It supports multiple scaling factors. Since multiple LoRA adapters may have + different scaling factors, we need multiple cos/sin caches. In this way, + instead of running rotary embedding kernel per lora, we can run multiple + lora in a batched way. + + In addition to that, we also keep the cos/sin cache for the scaling factor + of 1 (default) at all times. + + Exemplary for two scaling factors x=1, y and z with embeddings + [[x11, x12, ... x1m], ..., [xn1, xn2, ..., xnm]] and + [[y11, y12, ... y1o], ..., [yn1, yn2, ..., yno]], and + [[z11, z12, ... z1p], ..., [zn1, zn2, ..., znp]], + + we construct the cos/sin cache as follows: + [[x11, x12, ... x1m, y11, y12, ... y1o, z11, z12, ... z1p], + ... + [xn1, xn2, ... xnm, yn1, yn2, ... yno, zn1, zn2, ... znp]] + + We then use offsets to index into the cos/sin cache for + the respective scaling factors. + + The offset to cache can be accessed via `scaling_factor_to_offset` API. + + Credits to the Reddit user /u/kaiokendev + """ + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + scaling_factors: Union[List[float], float], + dtype: torch.dtype, + ) -> None: + if isinstance(scaling_factors, float): + scaling_factors = [scaling_factors] + self.scaling_factors: List[float] = scaling_factors # noqa + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + # Lazy initialized. + self._scaling_factor_to_offset: Dict[float, int] + + def _compute_cos_sin_cache(self) -> torch.Tensor: + inv_freq = self._compute_inv_freq(self.base) + cache_list: List[torch.Tensor] = [] + # offsets to the next cache in a tensor. + # Each offset corresponds to the same index in scaling_factors. + offsets: List[int] = [] + for scaling_factor in self.scaling_factors: + # NOTE(woosuk): self.max_position_embeddings is the original + # maximum length before applying the rope scaling. + # Thus, the maximum length after applying the rope scaling is + # self.max_position_embeddings * self.scaling_factor. + max_len = self.max_position_embeddings * scaling_factor + t = torch.arange(max_len, dtype=torch.float) + t = t / scaling_factor + + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() + sin = freqs.sin() + cache = torch.cat((cos, sin), dim=-1) + if not cache_list: + offset = 0 + else: + last_offset = offsets[-1] + next_max_len = cache_list[-1].shape[0] + offset = last_offset + next_max_len + offsets.append(offset) + cache_list.append(cache) + self._scaling_factor_to_offset = { + float(scaling_factor): offsets[i] + for i, scaling_factor in enumerate(self.scaling_factors) + } + assert len(self.scaling_factors) == len(offsets) + return torch.cat(cache_list, dim=0) + + @property + def scaling_factor_to_offset(self) -> Dict[float, int]: + return self._scaling_factor_to_offset + + +class DynamicNTKScalingRotaryEmbedding(RotaryEmbedding): + """RotaryEmbedding extended with Dynamic NTK scaling. + + Credits to the Reddit users /u/bloc97 and /u/emozilla + """ + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + scaling_factor: float, + dtype: torch.dtype, + ) -> None: + self.scaling_factor = scaling_factor + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + + def _compute_cos_sin_cache(self) -> torch.Tensor: + # NOTE(woosuk): self.max_position_embeddings is the original + # maximum length before applying the rope scaling. + # Thus, the maximum length after applying the rope scaling is + # self.max_position_embeddings * self.scaling_factor. + max_len = self.max_position_embeddings * self.scaling_factor + base = self.base * ( + (self.scaling_factor * max_len / self.max_position_embeddings) + - (self.scaling_factor - 1) + ) ** (self.rotary_dim / (self.rotary_dim - 2)) + inv_freq = self._compute_inv_freq(base) + t = torch.arange(max_len, 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 + + +# Inverse dim formula to find dim based on number of rotations +def _yarn_find_correction_dim( + num_rotations: int, + dim: int, + base: float = 10000, + max_position_embeddings: int = 2048, +) -> float: + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base) + ) + + +# Find dim range bounds based on rotations +def _yarn_find_correction_range( + low_rot: int, + high_rot: int, + dim: int, + base: float = 10000, + max_position_embeddings: int = 2048, +) -> Tuple[int, int]: + low = math.floor( + _yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings) + ) + high = math.ceil( + _yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings) + ) + return max(low, 0), min(high, dim - 1) # Clamp values just in case + + +def _yarn_linear_ramp_mask( + low: float, high: float, dim: int, dtype: torch.dtype, device: torch.device = None +) -> torch.Tensor: + if low == high: + high += 0.001 # Prevent singularity + + linear_func = (torch.arange(dim, dtype=dtype, device=device) - low) / (high - low) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func + + +def _yarn_get_mscale(scale: float = 1) -> float: + if scale <= 1: + return 1.0 + return 0.1 * math.log(scale) + 1.0 + + +class YaRNScalingRotaryEmbedding(RotaryEmbedding): + """RotaryEmbedding extended with YaRN method. + + Credits to Peng et al. github.com/jquesnelle/yarn + """ + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + scaling_factor: float, + dtype: torch.dtype, + *, + extrapolation_factor: float = 1, + attn_factor: float = 1, + beta_fast: int = 32, + beta_slow: int = 1, + ) -> None: + self.scaling_factor = scaling_factor + self.extrapolation_factor = extrapolation_factor + self.attn_factor = attn_factor + self.beta_fast = beta_fast + self.beta_slow = beta_slow + # Get n-d magnitude scaling corrected for interpolation + self.mscale = float(_yarn_get_mscale(self.scaling_factor) * attn_factor) + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + + def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor: + pos_freqs = self.base ** ( + torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim + ) + inv_freq_extrapolation = 1.0 / pos_freqs + inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs) + + low, high = _yarn_find_correction_range( + self.beta_fast, + self.beta_slow, + self.rotary_dim, + self.base, + self.max_position_embeddings, + ) + # Get n-d rotational scaling corrected for extrapolation + inv_freq_mask = ( + 1 + - _yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float) + ) * self.extrapolation_factor + inv_freq = ( + inv_freq_interpolation * (1 - inv_freq_mask) + + inv_freq_extrapolation * inv_freq_mask + ) + return inv_freq + + def _compute_cos_sin_cache(self) -> torch.Tensor: + inv_freq = self._compute_inv_freq(self.scaling_factor) + t = torch.arange( + self.max_position_embeddings * self.scaling_factor, dtype=torch.float32 + ) + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() * self.mscale + sin = freqs.sin() * self.mscale + cache = torch.cat((cos, sin), dim=-1) + return cache + + +class Phi3LongRoPEScaledRotaryEmbedding(nn.Module): + """Phi3 family of models scaled rotary embedding. + + Based on the original RotaryEmbedding implementation. + """ + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + original_max_position_embeddings: int, + base: int, + is_neox_style: bool, + dtype: torch.dtype, + short_factor: List[float], + long_factor: List[float], + short_mscale: Optional[float] = None, + long_mscale: Optional[float] = None, + ): + super().__init__() + + if is_neox_style is False: + raise ValueError( + "`Phi3LongRoPEScaledRotaryEmbedding` only supports neox_style." + ) + + self.rotary_dim = rotary_dim + self.head_size = head_size + self.max_position_embeddings = max_position_embeddings + self.original_max_position_embeddings = original_max_position_embeddings + self.base = base + self.short_factor = short_factor + self.long_factor = long_factor + + scale = self.max_position_embeddings / self.original_max_position_embeddings + if scale <= 1.0: + scaling_factor = 1.0 + else: + scaling_factor = math.sqrt( + 1 + math.log(scale) / math.log(self.original_max_position_embeddings) + ) + if short_mscale is None: + short_mscale = scaling_factor + if long_mscale is None: + long_mscale = scaling_factor + + self.short_mscale = short_mscale + self.long_mscale = long_mscale + + short_cache = self._compute_cos_sin_cache( + original_max_position_embeddings, short_factor, short_mscale + ) + short_cache = short_cache.to(dtype) + self.register_buffer("short_cos_sin_cache", short_cache, persistent=False) + + long_cache = self._compute_cos_sin_cache( + max_position_embeddings, long_factor, long_mscale + ) + long_cache = long_cache.to(dtype) + self.register_buffer("long_cos_sin_cache", long_cache, persistent=False) + + long_short_cache = torch.cat( + [self.short_cos_sin_cache, self.long_cos_sin_cache], dim=0 + ) + self.register_buffer( + "long_short_cos_sin_cache", long_short_cache, persistent=False + ) + + def _compute_inv_freq(self, rescale_factors: List[float]) -> torch.Tensor: + rescale_factors = torch.tensor(rescale_factors, dtype=torch.float32) + inv_freq = 1.0 / ( + rescale_factors + * ( + self.base + ** ( + torch.arange(0, self.rotary_dim, 2, dtype=torch.float) + / self.rotary_dim + ) + ) + ) + return inv_freq + + def _compute_cos_sin_cache( + self, + max_position_embeddings: int, + rescale_factors: List[float], + mscale: float, + ) -> torch.Tensor: + inv_freq = self._compute_inv_freq(rescale_factors) + t = torch.arange(max_position_embeddings, dtype=torch.float) + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() * mscale + sin = freqs.sin() * mscale + cache = torch.cat((cos, sin), dim=-1) + return cache + + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + query = query.view(*query.shape[:-1], -1, self.head_size) + key = key.view(*key.shape[:-1], -1, self.head_size) + + k = self.original_max_position_embeddings + long_prompt_offset = ( + torch.any(positions > k).float() * torch.full_like(positions, k) + ).long() + idx = ( + torch.add(positions, long_prompt_offset) + if long_prompt_offset is not None + else positions + ) + self.long_short_cos_sin_cache: torch.Tensor = self.long_short_cos_sin_cache.to( + idx.device + ) + idx = torch.add(idx, offsets) if offsets is not None else idx + cos_sin = torch.index_select(self.long_short_cos_sin_cache, 0, idx) + + cos, sin = cos_sin.chunk(2, dim=-1) + cos = cos.repeat(1, 2).unsqueeze(-2) + sin = sin.repeat(1, 2).unsqueeze(-2) + + query_rot = query[..., : self.rotary_dim] + query_pass = query[..., self.rotary_dim :] + query_rot = query_rot * cos + _rotate_neox(query_rot) * sin + query = torch.cat((query_rot, query_pass), dim=-1) + + key_rot = key[..., : self.rotary_dim] + key_pass = key[..., self.rotary_dim :] + key_rot = key_rot * cos + _rotate_neox(key_rot) * sin + key = torch.cat((key_rot, key_pass), dim=-1) + + return query.flatten(-2), key.flatten(-2) + + +def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class DeepseekScalingRotaryEmbedding(RotaryEmbedding): + """RotaryEmbedding extended with YaRN method. + + Credits to Peng et al. github.com/jquesnelle/yarn + """ + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + scaling_factor: float, + dtype: torch.dtype, + *, + extrapolation_factor: float = 1, + attn_factor: float = 1, + beta_fast: int = 32, + beta_slow: int = 1, + mscale: float = 1, + mscale_all_dim: float = 0, + device: Optional[str] = "cuda" if not _is_npu else "npu", + ) -> None: + self.scaling_factor = scaling_factor + self.extrapolation_factor = extrapolation_factor + self.attn_factor = attn_factor + self.beta_fast = beta_fast + self.beta_slow = beta_slow + # Get n-d magnitude scaling corrected for interpolation. + self.mscale = float( + yarn_get_mscale(self.scaling_factor, float(mscale)) + / yarn_get_mscale(self.scaling_factor, float(mscale_all_dim)) + * attn_factor + ) + self.device = device + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + + # Re-dispatch + if _is_hip: + self._forward_method = self.forward_native + + def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor: + pos_freqs = self.base ** ( + torch.arange(0, self.rotary_dim, 2, dtype=torch.float, device=self.device) + / self.rotary_dim + ) + inv_freq_extrapolation = 1.0 / pos_freqs + inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs) + + low, high = _yarn_find_correction_range( + self.beta_fast, + self.beta_slow, + self.rotary_dim, + self.base, + self.max_position_embeddings, + ) + # Get n-d rotational scaling corrected for extrapolation + inv_freq_mask = ( + 1 + - _yarn_linear_ramp_mask( + low, high, self.rotary_dim // 2, dtype=torch.float, device=self.device + ) + ) * self.extrapolation_factor + inv_freq = ( + inv_freq_interpolation * (1 - inv_freq_mask) + + inv_freq_extrapolation * inv_freq_mask + ) + return inv_freq + + def _compute_cos_sin_cache(self) -> torch.Tensor: + inv_freq = self._compute_inv_freq(self.scaling_factor) + t = torch.arange( + self.max_position_embeddings * self.scaling_factor, + device=self.device, + dtype=torch.float32, + ) + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() * self.mscale + sin = freqs.sin() * self.mscale + cache = torch.cat((cos, sin), dim=-1) + return cache + + def forward_native( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """PyTorch-native implementation equivalent to forward().""" + dtype = query.dtype + query_rot = query[..., : self.rotary_dim] + key_rot = key[..., : self.rotary_dim] + if self.rotary_dim < self.head_size: + query_pass = query[..., self.rotary_dim :] + key_pass = key[..., self.rotary_dim :] + + self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(positions.device) + cos_sin = self.cos_sin_cache[ + torch.add(positions, offsets) if offsets is not None else positions + ] + cos, sin = cos_sin.chunk(2, dim=-1) + if self.is_neox_style: + # NOTE(woosuk): Here we assume that the positions tensor has the + # shape [batch_size, seq_len]. + cos = cos.repeat(1, 1, 2).unsqueeze(-2) + sin = sin.repeat(1, 1, 2).unsqueeze(-2) + else: + cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) + sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) + + rotate_fn = _rotate_neox if self.is_neox_style else _rotate_gptj + query_rot = query_rot * cos + rotate_fn(query_rot) * sin + key_rot = key_rot * cos + rotate_fn(key_rot) * sin + + if self.rotary_dim < self.head_size: + query = torch.cat((query_rot, query_pass), dim=-1) + key = torch.cat((key_rot, key_pass), dim=-1) + else: + query = query_rot + key = key_rot + return query.to(dtype), key.to(dtype) + + def forward_npu( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # NOTE: now npu_mrope can only support `numQHeads*headSize <= 4096` pattern, + # and generalization to more scenarios will be supported in the future. + if query.shape[1] * query.shape[2] > 4096: + return self.forward_native(positions, query, key, offsets) + num_tokens = query.shape[0] + rotary_mode = "half" if self.is_neox_style else "interleave" + self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(positions.device) + query_rot = query[..., : self.rotary_dim] + key_rot = key[..., : self.rotary_dim] + if self.rotary_dim < self.head_size: + query_pass = query[..., self.rotary_dim :] + key_pass = key[..., self.rotary_dim :] + + query_rot, key_rot = torch_npu.npu_mrope( + torch.add(positions, offsets) if offsets is not None else positions, + query_rot.reshape(num_tokens, -1), + key_rot.reshape(num_tokens, -1), + self.cos_sin_cache, + self.rotary_dim, + mrope_section=[0, 0, 0], + rotary_mode=rotary_mode, + ) + query_rot = query_rot.reshape(num_tokens, -1, self.rotary_dim) + key_rot = key_rot.reshape(num_tokens, -1, self.rotary_dim) + + if self.rotary_dim < self.head_size: + query = torch.cat((query_rot, query_pass), dim=-1) + key = torch.cat((key_rot, key_pass), dim=-1) + else: + query = query_rot + key = key_rot + return query, key + + def forward_cpu( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + positions = torch.add(positions, offsets) if offsets is not None else positions + if _is_cpu_amx_available: + return torch.ops.sgl_kernel.rotary_embedding_cpu( + positions, query, key, self.head_size, self.cos_sin_cache, False + ) + else: + return self.forward_native(positions, query, key, offsets) + + +class Llama3RotaryEmbedding(RotaryEmbedding): + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + dtype: torch.dtype, + scaling_factor: float, + low_freq_factor: float, + high_freq_factor: float, + orig_max_position: int, + ) -> None: + self.scaling_factor = scaling_factor + self.low_freq_factor = low_freq_factor + self.high_freq_factor = high_freq_factor + self.orig_max_position = orig_max_position + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + + def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor: + inv_freqs = super()._compute_inv_freq(base) + low_freq_wavelen = self.orig_max_position / self.low_freq_factor + high_freq_wavelen = self.orig_max_position / self.high_freq_factor + + wave_len = 2 * math.pi / inv_freqs + if self.low_freq_factor != self.high_freq_factor: + smooth = (self.orig_max_position / wave_len - self.low_freq_factor) / ( + self.high_freq_factor - self.low_freq_factor + ) + else: + smooth = 0 + new_freqs = torch.where( + wave_len < high_freq_wavelen, + inv_freqs, + torch.where( + wave_len > low_freq_wavelen, + inv_freqs / self.scaling_factor, + (1 - smooth) * inv_freqs / self.scaling_factor + smooth * inv_freqs, + ), + ) + return new_freqs + + +class Llama4VisionRotaryEmbedding(RotaryEmbedding): + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + dtype: torch.dtype, + ): + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + + def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor: + inv_freqs = super()._compute_inv_freq(base) + inv_freqs = inv_freqs[: (self.rotary_dim // 2)] + return inv_freqs + + def _compute_cos_sin_cache(self) -> torch.Tensor: + inv_freq = self._compute_inv_freq(self.base) + + # self.max_position_embeddings here is number of image patches + # i.e. (image_size // patch_size) ** 2 + num_patches = self.max_position_embeddings + img_idx = torch.arange(num_patches, dtype=torch.int32).reshape(num_patches, 1) + img_idx = torch.cat([img_idx, img_idx[:1]], dim=0) + img_idx[-1, -1] = -2 # set to ID_CLS_TOKEN + num_patches_single_dim = int(math.sqrt(num_patches)) + frequencies_x = img_idx % num_patches_single_dim + frequencies_y = img_idx // num_patches_single_dim + freqs_x = ( + (frequencies_x + 1)[..., None] * inv_freq[None, None, :] + ).repeat_interleave(2, dim=-1) + freqs_y = ( + (frequencies_y + 1)[..., None] * inv_freq[None, None, :] + ).repeat_interleave(2, dim=-1) + freqs = torch.cat([freqs_x, freqs_y], dim=-1).float().contiguous()[..., ::2] + freqs = freqs.masked_fill(img_idx.reshape(-1, 1, 1) < 0, 0) + cache = torch.view_as_complex( + torch.stack([torch.cos(freqs), torch.sin(freqs)], dim=-1) + ) + return cache + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device) + query_ = torch.view_as_complex(query.float().reshape(*query.shape[:-1], -1, 2)) + key_ = torch.view_as_complex(key.float().reshape(*key.shape[:-1], -1, 2)) + broadcast_shape = [ + d if i == 1 or i == (query_.ndim - 1) else 1 + for i, d in enumerate(query_.shape) + ] + freqs_ci = self.cos_sin_cache.view(*broadcast_shape) + query_out = torch.view_as_real(query_ * freqs_ci).flatten(3) + key_out = torch.view_as_real(key_ * freqs_ci).flatten(3) + return query_out.type_as(query), key_out.type_as(key) + + +class DynamicNTKAlphaRotaryEmbedding(RotaryEmbedding): + """RotaryEmbedding extended with Dynamic NTK scaling. + + Credits to the Reddit users /u/bloc97 and /u/emozilla + """ + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + scaling_alpha: float, + dtype: torch.dtype, + ) -> None: + self.scaling_alpha = scaling_alpha + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + + def _compute_cos_sin_cache(self) -> torch.Tensor: + max_len = self.max_position_embeddings + base = self.base * self.scaling_alpha ** ( + self.rotary_dim / (self.rotary_dim - 2) + ) + + inv_freq = self._compute_inv_freq(base) + t = torch.arange(max_len, 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 + + +class MRotaryEmbedding(RotaryEmbedding): + """Rotary Embedding with Multimodal Sections.""" + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + dtype: torch.dtype, + mrope_section: Optional[List[int]] = None, + ) -> None: + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + + self.mrope_section = mrope_section + if self.mrope_section: + expected_sum = rotary_dim // 2 + actual_sum = sum(self.mrope_section) + if actual_sum != expected_sum: + print( + f"MRoPE section sum mismatch: expected {expected_sum}, got {actual_sum}. " + f"Adjusting mrope_section to match rotary_dim // 2 = {expected_sum}" + ) + # Auto-correct by scaling the mrope_section proportionally + if actual_sum > 0: + scale_factor = expected_sum / actual_sum + self.mrope_section = [ + max(1, int(section * scale_factor)) + for section in self.mrope_section + ] + # Ensure the sum exactly matches by adjusting the last element + current_sum = sum(self.mrope_section) + if current_sum != expected_sum: + self.mrope_section[-1] += expected_sum - current_sum + else: + # If all sections are 0, create a default distribution + self.mrope_section = [ + expected_sum // len(self.mrope_section) + ] * len(self.mrope_section) + # Handle remainder + remainder = expected_sum % len(self.mrope_section) + for i in range(remainder): + self.mrope_section[i] += 1 + + print( + f"Corrected mrope_section: {self.mrope_section} (sum={sum(self.mrope_section)})" + ) + + @torch.compile(dynamic=True) + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """PyTorch-native implementation equivalent to forward(). + + Args: + positions: + [num_tokens,] (text only) or + [3, num_tokens] (T/H/W positions with multimodal inputs) + query: [num_tokens, num_heads * head_size] + key: [num_tokens, num_kv_heads * head_size] + """ + assert positions.ndim == 1 or positions.ndim == 2 + + num_tokens = positions.shape[-1] + cos_sin = self.cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + if positions.ndim == 2: + assert self.mrope_section + + cos = torch.cat( + [m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))], + dim=-1, + ) + sin = torch.cat( + [m[i] for i, m in enumerate(sin.split(self.mrope_section, dim=-1))], + dim=-1, + ) + + query_shape = query.shape + query = query.view(num_tokens, -1, self.head_size) + query_rot = query[..., : self.rotary_dim] + query_pass = query[..., self.rotary_dim :] + query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style) + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) + + key_shape = key.shape + key = key.view(num_tokens, -1, self.head_size) + key_rot = key[..., : self.rotary_dim] + key_pass = key[..., self.rotary_dim :] + key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style) + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) + return query, key + + # Copied from https://github.com/huggingface/transformers/blob/c8e0e603de9b3d49161a15fe6e8ea84badfb5d02/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1439 + @staticmethod + def get_rope_index( + spatial_merge_size: int, + image_token_id: int, + video_token_id: int, + vision_start_token_id: int, + model_type: str, + tokens_per_second: Optional[int] = None, + input_ids: Optional[torch.LongTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + mrope_position_deltas = [] + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): + total_input_ids = input_ids + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + for i, input_ids in enumerate(total_input_ids): + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere( + input_ids == vision_start_token_id + ).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + second_per_grid_t = 0 + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + if second_per_grid_ts is not None: + second_per_grid_t = second_per_grid_ts[video_index] + else: + second_per_grid_t = 1.0 + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + if model_type == "qwen2_5_vl": + range_tensor = torch.arange(llm_grid_t).view(-1, 1) + expanded_range = range_tensor.expand( + -1, llm_grid_h * llm_grid_w + ) + + time_tensor = ( + expanded_range * second_per_grid_t * tokens_per_second + ) + + time_tensor_long = time_tensor.long() + t_index = time_tensor_long.flatten() + elif model_type == "qwen2_vl": + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + else: + raise RuntimeError("Unimplemented") + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, :] = llm_positions.to(position_ids.device) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + s = input_ids.shape[1] + position_ids = torch.arange(s) + position_ids = ( + position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max( + -1, keepdim=True + )[0] + mrope_position_deltas = max_position_ids + 1 - s + return position_ids, mrope_position_deltas + + # Adapted from https://github.com/vllm-project/vllm/blob/3779eb8c81449b924a23457fc77e45a0e6171178/vllm/model_executor/layers/rotary_embedding.py#L1120 + @staticmethod + def get_rope_index_glm4v( + input_ids: torch.Tensor, + hf_config: Any, + image_grid_thw: Union[list[list[int]], torch.Tensor], + video_grid_thw: Union[list[list[int]], torch.Tensor], + attention_mask: torch.Tensor, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Get mrope input positions and delta value for GLM4V.""" + image_token_id = hf_config.image_token_id + video_start_token_id = hf_config.video_start_token_id + video_end_token_id = hf_config.video_end_token_id + spatial_merge_size = hf_config.vision_config.spatial_merge_size + + mrope_position_deltas = [] + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + video_group_index = 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + input_tokens = input_ids.tolist() + + input_token_type = [] + video_check_flg = False + for token in input_tokens: + if token == video_start_token_id: + video_check_flg = True + elif token == video_end_token_id: + video_check_flg = False + + if token == image_token_id and not video_check_flg: + input_token_type.append("image") + elif token == image_token_id and video_check_flg: + input_token_type.append("video") + else: + input_token_type.append("text") + + input_type_group = [] + for key, group in itertools.groupby( + enumerate(input_token_type), lambda x: x[1] + ): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + llm_pos_ids_list = [] + video_frame_num = 1 + for modality_type, start_idx, end_idx in input_type_group: + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + + if modality_type == "image": + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + st_idx + ) + + image_index += 1 + video_frame_num = 1 + + elif modality_type == "video": + t, h, w = ( + video_frame_num, + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + + llm_grid_t, llm_grid_h, llm_grid_w = ( + t, + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + + for t_idx in range(llm_grid_t): + t_index = ( + torch.tensor(t_idx) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(1, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(1, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + st_idx + ) + + video_group_index += 1 + + if video_group_index >= video_grid_thw[video_index][0]: + video_index += 1 + video_group_index = 0 + + video_frame_num += 1 + + else: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + video_frame_num = 1 + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = ( + position_ids.unsqueeze(0) + .expand(3, -1, -1) + .to(attention_mask.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max( + -1, keepdim=True + )[0] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + + return position_ids, mrope_position_deltas + + +class DualChunkRotaryEmbedding(CustomOp): + """Rotary positional embedding for Dual Chunk Attention.""" + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + dtype: torch.dtype, + chunk_size: int, + local_size: int, + ) -> 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.chunk_size = chunk_size + self.local_size = local_size + self.dtype = dtype + self.device = torch.device(f"cuda:{torch.cuda.current_device()}") + (q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache) = ( + self._compute_cos_sin_cache() + ) + + self.register_buffer("cos_sin_q_cache", q_cache, persistent=False) + self.register_buffer("cos_sin_qc_cache", qc_cache, persistent=False) + self.register_buffer("cos_sin_k_cache", k_cache, persistent=False) + self.register_buffer( + "cos_sin_qc_no_clamp_cache", qc_no_clamp_cache, persistent=False + ) + self.register_buffer("cos_sin_q_inter_cache", q_inter_cache, persistent=False) + + def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor: + """Compute the inverse frequency.""" + # NOTE(woosuk): The HF implementation uses `torch.arange(...).float()`. + # However, we use `torch.arange(..., dtype=torch.float)` instead to + # avoid numerical issues with large base values (e.g., 10000000). + # This may cause a slight numerical difference between the HF + # implementation and ours. + # NOTE(woosuk): To exactly match the HF implementation, we need to + # use CPU to compute the cache and then move it to GPU. However, we + # create the cache on GPU for faster initialization. This may cause + # a slight numerical difference between the HF implementation and ours. + 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) + chunk_len = self.chunk_size - self.local_size + q_t = torch.arange(chunk_len, dtype=torch.float) + qc_t = (torch.arange(chunk_len, dtype=torch.float) + chunk_len).clamp( + max=self.chunk_size + ) + k_t = torch.arange(self.max_position_embeddings, dtype=torch.float) % chunk_len + + # count from chunk_len, no clamp(self.chunk_size) restriction + qc_no_clamp_t = torch.arange(chunk_len, dtype=torch.float) + chunk_len + # count from self.chunk_size for q_inter's rope + q_inter_t = torch.arange(chunk_len, dtype=torch.float) + self.chunk_size + + q_freqs = torch.outer(q_t, inv_freq) + qc_freqs = torch.outer(qc_t, inv_freq) + k_freqs = torch.outer(k_t, inv_freq) + qc_no_clamp_freqs = torch.outer(qc_no_clamp_t, inv_freq) + q_inter_freqs = torch.outer(q_inter_t, inv_freq) + + q_cos = q_freqs.cos() + q_sin = q_freqs.sin() + qc_cos = qc_freqs.cos() + qc_sin = qc_freqs.sin() + k_cos = k_freqs.cos() + k_sin = k_freqs.sin() + + qc_no_clamp_cos = qc_no_clamp_freqs.cos() + qc_no_clamp_sin = qc_no_clamp_freqs.sin() + q_inter_cos = q_inter_freqs.cos() + q_inter_sin = q_inter_freqs.sin() + + q_cache = torch.cat((q_cos, q_sin), dim=-1).to( + dtype=self.dtype, device=self.device + ) + qc_cache = torch.cat((qc_cos, qc_sin), dim=-1).to( + dtype=self.dtype, device=self.device + ) + k_cache = torch.cat((k_cos, k_sin), dim=-1).to( + dtype=self.dtype, device=self.device + ) + qc_no_clamp_cache = torch.cat((qc_no_clamp_cos, qc_no_clamp_sin), dim=-1).to( + dtype=self.dtype, device=self.device + ) + q_inter_cache = torch.cat((q_inter_cos, q_inter_sin), dim=-1).to( + dtype=self.dtype, device=self.device + ) + return q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache + + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + query = query.view(*query.shape[:-1], -1, self.head_size) + key = key.view(*key.shape[:-1], -1, self.head_size) + query_rot = query[..., : self.rotary_dim] + key_rot = key[..., : self.rotary_dim] + if self.rotary_dim < self.head_size: + query_pass = query[..., self.rotary_dim :] + key_pass = key[..., self.rotary_dim :] + else: + query_pass = None + key_pass = None + + positions_with_offsets = ( + torch.add(positions, offsets) if offsets is not None else positions + ) + key = self._apply_rotary_embedding( + self.cos_sin_k_cache[positions_with_offsets], key_rot, key_pass + ) + chunk_len = self.chunk_size - self.local_size + query = self._apply_rotary_embedding( + self.cos_sin_q_cache[positions_with_offsets % chunk_len], + query_rot, + query_pass, + ) + query_succ = self._apply_rotary_embedding( + self.cos_sin_qc_cache[positions_with_offsets % chunk_len], + query_rot, + query_pass, + ) + query_inter = self._apply_rotary_embedding( + self.cos_sin_qc_cache[chunk_len - 1].repeat(positions.shape[0], 1), + query_rot, + query_pass, + ) + query_succ_critical = self._apply_rotary_embedding( + self.cos_sin_qc_no_clamp_cache[positions_with_offsets % chunk_len], + query_rot, + query_pass, + ) + query_inter_critical = self._apply_rotary_embedding( + self.cos_sin_q_inter_cache[positions_with_offsets % chunk_len], + query_rot, + query_pass, + ) + + # merge query into one tensor to simplify the interfaces + query = torch.cat( + ( + query, + query_succ, + query_inter, + query_succ_critical, + query_inter_critical, + ), + dim=-1, + ) + return query, key + + def _apply_rotary_embedding(self, cos_sin, hidden_rot, hidden_pass): + cos, sin = cos_sin.chunk(2, dim=-1) + if self.is_neox_style: + # NOTE(woosuk): Here we assume that the positions tensor has the + # shape [batch_size, seq_len]. + cos = cos.repeat(1, 1, 2).unsqueeze(-2) + sin = sin.repeat(1, 1, 2).unsqueeze(-2) + else: + cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) + sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) + rotate_fn = _rotate_neox if self.is_neox_style else _rotate_gptj + hidden_rot = hidden_rot * cos + rotate_fn(hidden_rot) * sin + + if self.rotary_dim < self.head_size: + hidden = torch.cat((hidden_rot, hidden_pass), dim=-1) + else: + hidden = hidden_rot + return hidden.flatten(-2).squeeze(0) + + def extra_repr(self) -> str: + s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" + s += f", max_position_embeddings={self.max_position_embeddings}" + s += f", base={self.base}, is_neox_style={self.is_neox_style}" + s += f", chunk_size={self.chunk_size}, local_size={self.local_size}" + return s + + +_ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {} + + +def get_rope( + head_size: int, + rotary_dim: int, + max_position: int, + base: int, + is_neox_style: bool = True, + rope_scaling: Optional[Dict[str, Any]] = None, + dtype: Optional[torch.dtype] = None, + partial_rotary_factor: float = 1.0, + dual_chunk_attention_config: Optional[Dict[str, Any]] = None, +) -> RotaryEmbedding: + if dtype is None: + dtype = torch.get_default_dtype() + if rope_scaling is not None: + # Transforms every value that is a list into a tuple for caching calls + rope_scaling_tuple = { + k: tuple(v) if isinstance(v, list) else v for k, v in rope_scaling.items() + } + rope_scaling_args = tuple(rope_scaling_tuple.items()) + else: + rope_scaling_args = None + + if dual_chunk_attention_config is not None: + dual_chunk_attention_tuple = { + k: tuple(v) if isinstance(v, list) else v + for k, v in dual_chunk_attention_config.items() + if k != "sparse_attention_config" + } + dual_chunk_attention_args = tuple(dual_chunk_attention_tuple.items()) + else: + dual_chunk_attention_args = None + + if partial_rotary_factor < 1.0: + rotary_dim = int(rotary_dim * partial_rotary_factor) + key = ( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + rope_scaling_args, + dual_chunk_attention_args, + dtype, + ) + if key in _ROPE_DICT: + return _ROPE_DICT[key] + + if dual_chunk_attention_config is not None: + extra_kwargs = { + k: v + for k, v in dual_chunk_attention_config.items() + if k in ("chunk_size", "local_size") + } + rotary_emb = DualChunkRotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + dtype, + **extra_kwargs, + ) + elif rope_scaling is None: + rotary_emb = RotaryEmbedding( + head_size, rotary_dim, max_position, base, is_neox_style, dtype + ) + else: + if "rope_type" in rope_scaling: + scaling_type = rope_scaling["rope_type"] + elif "type" in rope_scaling: + scaling_type = rope_scaling["type"] + else: + raise ValueError("Unknown RoPE scaling type") + + if scaling_type == "llama3": + scaling_factor = rope_scaling["factor"] + low_freq_factor = rope_scaling["low_freq_factor"] + high_freq_factor = rope_scaling["high_freq_factor"] + original_max_position = rope_scaling["original_max_position_embeddings"] + rotary_emb = Llama3RotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + dtype, + scaling_factor, + low_freq_factor, + high_freq_factor, + original_max_position, + ) + elif scaling_type == "default": + if "mrope_section" in rope_scaling: + rotary_emb = MRotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + dtype, + mrope_section=rope_scaling["mrope_section"], + ) + else: + rotary_emb = RotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + dtype, + ) + elif scaling_type == "linear": + scaling_factor = rope_scaling["factor"] + rotary_emb = LinearScalingRotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + scaling_factor, + dtype, + ) + elif scaling_type == "dynamic": + scaling_factor = rope_scaling["factor"] + if "alpha" in rope_scaling: + rotary_emb = DynamicNTKAlphaRotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + rope_scaling["alpha"], + dtype, + ) + else: + rotary_emb = DynamicNTKScalingRotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + scaling_factor, + dtype, + ) + elif scaling_type == "yarn": + scaling_factor = rope_scaling["factor"] + original_max_position = rope_scaling["original_max_position_embeddings"] + extra_kwargs = { + k: v + for k, v in rope_scaling.items() + if k + in ("extrapolation_factor", "attn_factor", "beta_fast", "beta_slow") + } + rotary_emb = YaRNScalingRotaryEmbedding( + head_size, + rotary_dim, + original_max_position, + base, + is_neox_style, + scaling_factor, + dtype, + **extra_kwargs, + ) + elif scaling_type == "deepseek_yarn": + scaling_factor = rope_scaling["factor"] + original_max_position = rope_scaling["original_max_position_embeddings"] + # assert max_position == original_max_position * scaling_factor + extra_kwargs = { + k: v + for k, v in rope_scaling.items() + if k + in ( + "extrapolation_factor", + "attn_factor", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ) + } + rotary_emb = DeepseekScalingRotaryEmbedding( + head_size, + rotary_dim, + original_max_position, + base, + is_neox_style, + scaling_factor, + dtype, + **extra_kwargs, + ) + elif scaling_type == "longrope": + short_factor = rope_scaling["short_factor"] + long_factor = rope_scaling["long_factor"] + original_max_position = rope_scaling["original_max_position_embeddings"] + extra_kwargs = { + k: v + for k, v in rope_scaling.items() + if k in ("short_mscale", "long_mscale") + } + rotary_emb = Phi3LongRoPEScaledRotaryEmbedding( + head_size, + rotary_dim, + max_position, + original_max_position, + base, + is_neox_style, + dtype, + short_factor, + long_factor, + **extra_kwargs, + ) + else: + raise ValueError(f"Unknown RoPE scaling type {scaling_type}") + _ROPE_DICT[key] = rotary_emb + return rotary_emb + + +# Copied from transformers +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb_native( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + unsqueeze_dim=1, +) -> Tuple[torch.Tensor, torch.Tensor]: + orig_q_dtype = q.dtype + orig_k_dtype = k.dtype + q, k = q.float(), k.float() + + # embedding is performed in float + cos = cos.unsqueeze(unsqueeze_dim).float() + sin = sin.unsqueeze(unsqueeze_dim).float() + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + + q_embed = q_embed.to(orig_q_dtype) + k_embed = k_embed.to(orig_k_dtype) + + return q_embed, k_embed + + +def apply_rotary_pos_emb_npu( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + unsqueeze_dim=1, +) -> Tuple[torch.Tensor, torch.Tensor]: + if q.shape[1] != 128: + return apply_rotary_pos_emb_native(q, k, cos, sin, unsqueeze_dim) + cos = cos.unsqueeze(unsqueeze_dim) + cos = torch.transpose(cos, 1, 2) + sin = sin.unsqueeze(unsqueeze_dim) + sin = torch.transpose(sin, 1, 2) + q = torch.transpose(q, 1, 2) + k = torch.transpose(k, 1, 2) + q_embed, k_embed = torch_npu.npu_apply_rotary_pos_emb(q, k, cos, sin) + q_embed = torch.transpose(q_embed, 1, 2) + k_embed = torch.transpose(k_embed, 1, 2) + return q_embed, k_embed + + +if _is_npu: + apply_rotary_pos_emb = apply_rotary_pos_emb_npu +else: + apply_rotary_pos_emb = apply_rotary_pos_emb_native + + +def get_rope_cpu( + head_size: int, + rotary_dim: int, + max_position: int, + base: int, + is_neox_style: bool = True, + rope_scaling: Optional[Dict[str, Any]] = None, + dtype: Optional[torch.dtype] = None, + partial_rotary_factor: float = 1.0, + device: Optional[str] = None, +) -> RotaryEmbedding: + if dtype is None: + dtype = torch.get_default_dtype() + if rope_scaling is not None: + # Transforms every value that is a list into a tuple for caching calls + rope_scaling_tuple = { + k: tuple(v) if isinstance(v, list) else v for k, v in rope_scaling.items() + } + rope_scaling_args = tuple(rope_scaling_tuple.items()) + else: + rope_scaling_args = None + if partial_rotary_factor < 1.0: + rotary_dim = int(rotary_dim * partial_rotary_factor) + key = ( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + rope_scaling_args, + dtype, + ) + if key in _ROPE_DICT: + return _ROPE_DICT[key] + + assert rope_scaling is not None + scaling_type = rope_scaling["rope_type"] + assert ( + scaling_type == "deepseek_yarn" + ), "Only deepseek_yarn is supported for CPU for now" + + scaling_factor = rope_scaling["factor"] + original_max_position = rope_scaling["original_max_position_embeddings"] + extra_kwargs = { + k: v + for k, v in rope_scaling.items() + if k + in ( + "extrapolation_factor", + "attn_factor", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ) + } + extra_kwargs["device"] = device + rotary_emb = DeepseekScalingRotaryEmbedding( + head_size, + rotary_dim, + original_max_position, + base, + is_neox_style, + scaling_factor, + dtype, + **extra_kwargs, + ) + + _ROPE_DICT[key] = rotary_emb + return rotary_emb + + +def get_rope_wrapper( + head_size: int, + rotary_dim: int, + max_position: int, + base: int, + is_neox_style: bool = True, + rope_scaling: Optional[Dict[str, Any]] = None, + dtype: Optional[torch.dtype] = None, + partial_rotary_factor: float = 1.0, + device: Optional[str] = None, +): + if device != "cpu": + wrapper = aiter_get_rope if _use_aiter else get_rope + return wrapper( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + rope_scaling, + dtype, + partial_rotary_factor, + ) + + return get_rope_cpu( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + rope_scaling, + dtype, + partial_rotary_factor, + device, + ) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py new file mode 100644 index 000000000..56a831f2d --- /dev/null +++ b/python/sglang/srt/layers/sampler.py @@ -0,0 +1,294 @@ +import logging +from typing import List + +import torch +import torch.distributed as dist +from torch import nn + +from sglang.srt.distributed import get_tp_group +from sglang.srt.layers.dp_attention import ( + get_attention_tp_group, + is_dp_attention_enabled, +) +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo +from sglang.srt.utils import crash_on_warnings, get_bool_env_var, is_cuda + +if is_cuda(): + from sgl_kernel import ( + min_p_sampling_from_probs, + top_k_renorm_prob, + top_k_top_p_sampling_from_probs, + top_p_renorm_prob, + ) + + +logger = logging.getLogger(__name__) + +SYNC_TOKEN_IDS_ACROSS_TP = get_bool_env_var("SYNC_TOKEN_IDS_ACROSS_TP") +RETURN_ORIGINAL_LOGPROB = get_bool_env_var("RETURN_ORIGINAL_LOGPROB") + + +class Sampler(nn.Module): + def __init__(self): + super().__init__() + self.use_nan_detection = global_server_args_dict["enable_nan_detection"] + self.tp_sync_group = get_tp_group().device_group + + if is_dp_attention_enabled(): + self.tp_sync_group = get_attention_tp_group().device_group + + def forward( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + return_logprob: bool, + top_logprobs_nums: List[int], + token_ids_logprobs: List[List[int]], + ): + """Run a sampler & compute logprobs and update logits_output accordingly. + + Args: + logits_output: The logits from the model forward + sampling_info: Metadata for sampling + return_logprob: If set, store the output logprob information to + logits_output + top_logprobs_nums: Number of top lobprobs per sequence in a batch + batch_next_token_ids: next token IDs. If set, skip sampling and only + compute output logprobs It is used for speculative decoding which + performs sampling in draft workers. + """ + logits = logits_output.next_token_logits + + # Apply the custom logit processors if registered in the sampling info. + if sampling_info.has_custom_logit_processor: + apply_custom_logit_processor(logits, sampling_info) + + if self.use_nan_detection and torch.any(torch.isnan(logits)): + logger.warning("Detected errors during sampling! NaN in the logits.") + logits = torch.where( + torch.isnan(logits), torch.full_like(logits, -1e5), logits + ) + if crash_on_warnings(): + raise ValueError("Detected errors during sampling! NaN in the logits.") + + if sampling_info.is_all_greedy: + # Use torch.argmax if all requests use greedy sampling + batch_next_token_ids = torch.argmax(logits, -1) + if return_logprob: + logprobs = torch.nn.functional.log_softmax(logits, dim=-1) + + else: + # Post process original logits. if temperatures are all 1.0, no need to rescale + if return_logprob and RETURN_ORIGINAL_LOGPROB: + logprobs = torch.softmax(logits, dim=-1) + + # Post process logits + logits.div_(sampling_info.temperatures) + logits[:] = torch.softmax(logits, dim=-1) + probs = logits + del logits + + if True: # Keep this redundant check to simplify some internal code sync + if global_server_args_dict["sampling_backend"] == "flashinfer": + if sampling_info.need_min_p_sampling: + probs = top_k_renorm_prob(probs, sampling_info.top_ks) + probs = top_p_renorm_prob(probs, sampling_info.top_ps) + batch_next_token_ids = min_p_sampling_from_probs( + probs, sampling_info.min_ps + ) + else: + batch_next_token_ids = top_k_top_p_sampling_from_probs( + probs.contiguous(), + sampling_info.top_ks, + sampling_info.top_ps, + filter_apply_order="joint", + check_nan=self.use_nan_detection, + ) + elif global_server_args_dict["sampling_backend"] == "pytorch": + # A slower fallback implementation with torch native operations. + batch_next_token_ids = top_k_top_p_min_p_sampling_from_probs_torch( + probs, + sampling_info.top_ks, + sampling_info.top_ps, + sampling_info.min_ps, + sampling_info.need_min_p_sampling, + ) + else: + raise ValueError( + f"Invalid sampling backend: {global_server_args_dict['sampling_backend']}" + ) + + if return_logprob: + # clamp to avoid -inf + if RETURN_ORIGINAL_LOGPROB: + logprobs = torch.log(logprobs).clamp( + min=torch.finfo(logprobs.dtype).min + ) + else: + logprobs = torch.log(probs).clamp(min=torch.finfo(probs.dtype).min) + + # Attach logprobs to logits_output (in-place modification) + if return_logprob: + if any(x > 0 for x in top_logprobs_nums): + ( + logits_output.next_token_top_logprobs_val, + logits_output.next_token_top_logprobs_idx, + ) = get_top_logprobs(logprobs, top_logprobs_nums) + + if any(x is not None for x in token_ids_logprobs): + ( + logits_output.next_token_token_ids_logprobs_val, + logits_output.next_token_token_ids_logprobs_idx, + ) = get_token_ids_logprobs(logprobs, token_ids_logprobs) + + logits_output.next_token_logprobs = logprobs[ + torch.arange(len(batch_next_token_ids), device=sampling_info.device), + batch_next_token_ids, + ] + + if SYNC_TOKEN_IDS_ACROSS_TP or sampling_info.grammars: + # For performance reasons, SGLang does not sync the final token IDs across TP ranks by default. + # This saves one all-reduce, but the correctness of this approach depends on the determinism of several operators: + # the last all-reduce, the last lm_head matmul, and all sampling kernels. + # These kernels are deterministic in most cases, but there are some rare instances where they are not deterministic. + # In such cases, enable this env variable to prevent hanging due to TP ranks becoming desynchronized. + # When using xgrammar, this becomes more likely so we also do the sync when grammar is used. + + torch.distributed.all_reduce( + batch_next_token_ids, + op=dist.ReduceOp.MIN, + group=self.tp_sync_group, + ) + + return batch_next_token_ids + + +def top_k_top_p_min_p_sampling_from_probs_torch( + probs: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: torch.Tensor, + need_min_p_sampling: bool, +): + """A top-k, top-p and min-p sampling implementation with native pytorch operations.""" + probs_sort, probs_idx = probs.sort(dim=-1, descending=True) + probs_sum = torch.cumsum(probs_sort, dim=-1) + probs_sort[ + torch.arange(0, probs.shape[-1], device=probs.device).view(1, -1) + >= top_ks.view(-1, 1) + ] = 0.0 + probs_sort[(probs_sum - probs_sort) > top_ps.view(-1, 1)] = 0.0 + + if need_min_p_sampling: + min_p_thresholds = probs_sort[:, 0] * min_ps + probs_sort[probs_sort < min_p_thresholds.view(-1, 1)] = 0.0 + + sampled_index = torch.multinomial(probs_sort, num_samples=1) + # int32 range is enough to represent the token ids + probs_idx = probs_idx.to(torch.int32) + batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(-1) + return batch_next_token_ids + + +def sampling_from_probs_torch(probs: torch.Tensor): + """A sampling implementation with native pytorch operations, without + top-k, top-p, or min-p filtering.""" + sampled_index = torch.multinomial(probs, num_samples=1) + batch_next_token_ids = sampled_index.view(-1).to(torch.int32) + return batch_next_token_ids + + +def top_p_normalize_probs_torch( + probs: torch.Tensor, + top_ps: torch.Tensor, +): + # See also top_k_top_p_min_p_sampling_from_probs_torch + probs_sort, probs_idx = probs.sort(dim=-1, descending=True) + probs_sum = torch.cumsum(probs_sort, dim=-1) + probs_sort[(probs_sum - probs_sort) > top_ps.view(-1, 1)] = 0.0 + probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True)) + return torch.zeros_like(probs_sort).scatter_(-1, probs_idx, probs_sort) + + +def get_top_logprobs( + logprobs: torch.Tensor, + top_logprobs_nums: List[int], +): + max_k = max(top_logprobs_nums) + ret = logprobs.topk(max_k, dim=1) + values = ret.values.tolist() + indices = ret.indices.tolist() + + output_top_logprobs_val = [] + output_top_logprobs_idx = [] + for i, k in enumerate(top_logprobs_nums): + output_top_logprobs_val.append(values[i][:k]) + output_top_logprobs_idx.append(indices[i][:k]) + + return ( + output_top_logprobs_val, + output_top_logprobs_idx, + ) + + +def get_token_ids_logprobs( + logprobs: torch.Tensor, + token_ids_logprobs: List[List[int]], +): + output_token_ids_logprobs_val = [] + output_token_ids_logprobs_idx = [] + for i, token_ids in enumerate(token_ids_logprobs): + if token_ids is not None: + output_token_ids_logprobs_val.append(logprobs[i, token_ids].tolist()) + output_token_ids_logprobs_idx.append(token_ids) + else: + output_token_ids_logprobs_val.append([]) + output_token_ids_logprobs_idx.append([]) + + return ( + output_token_ids_logprobs_val, + output_token_ids_logprobs_idx, + ) + + +def apply_custom_logit_processor( + logits: torch.Tensor, + sampling_batch_info: SamplingBatchInfo, + num_tokens_in_batch: int = 1, +): + """Apply custom logit processors to the logits. + This function will modify the logits in-place. + num_tokens_in_batch is needed to support spec decoding, where each batch can contain multiple + tokens. By default, we assume each batch contains only 1 token. + """ + + assert logits.shape[0] == len(sampling_batch_info) * num_tokens_in_batch, ( + f"The batch size of logits ({logits.shape[0]}) does not match the batch size of " + f"sampling_batch_info ({len(sampling_batch_info)}) x num_tokens_in_batch " + f"({num_tokens_in_batch})" + ) + + for _, ( + processor, + batch_mask, + ) in sampling_batch_info.custom_logit_processor.items(): + # Get the batch indices that need to be processed + batch_indices = batch_mask.nonzero(as_tuple=True)[0] + + assert batch_mask.shape[0] == len(sampling_batch_info), ( + f"The number of batch mask ({batch_mask.shape[0]}) does not match the number of " + f"sampling_batch_info ({len(sampling_batch_info)})" + ) + batch_mask = torch.repeat_interleave(batch_mask, num_tokens_in_batch) + + # Apply the processor to the logits + logits[batch_mask] = processor( + logits[batch_mask], + [sampling_batch_info.custom_params[i] for i in batch_indices], + ) + + logger.debug( + f"Custom logit processor {processor.__class__.__name__} is applied." + ) diff --git a/python/sglang/srt/layers/torchao_utils.py b/python/sglang/srt/layers/torchao_utils.py new file mode 100644 index 000000000..e08abd5ae --- /dev/null +++ b/python/sglang/srt/layers/torchao_utils.py @@ -0,0 +1,124 @@ +""" +Common utilities for torchao. +""" + +import logging +import os +import pwd +from typing import Callable, Optional + +import torch + +logger = logging.getLogger(__name__) + + +def get_gemlite_cache_path() -> str: + return f"/tmp/{pwd.getpwuid(os.getuid()).pw_gecos}_gemlite.json" + + +def save_gemlite_cache(print_error: bool = False) -> bool: + try: + from gemlite.core import GemLiteLinearTriton + + GemLiteLinearTriton.cache_config(get_gemlite_cache_path()) + except Exception: + if print_error: + logger.error("Failed to save the GemLite cache.") + return False + return True + + +def proj_filter( + module: torch.nn.Module, + fqn: str, +): + """Filter function for quantizing projection layers.""" + return "proj" in fqn + + +def apply_torchao_config_to_model( + model: torch.nn.Module, + torchao_config: str, + filter_fn: Optional[Callable] = proj_filter, +): + """Quantize a modelwith torchao quantization specified by torchao_config + + Args: + `model`: a model to be quantized based on torchao_config + `torchao_config` (str): type of quantization and their arguments we want to use to + quantize the model, e.g. int4wo-128 means int4 weight only quantization with group_size + 128 + """ + # Lazy import to suppress some warnings + from torchao.quantization import ( + float8_dynamic_activation_float8_weight, + float8_weight_only, + int4_weight_only, + int8_dynamic_activation_int8_weight, + int8_weight_only, + quantize_, + ) + from torchao.quantization.observer import PerRow, PerTensor + + if torchao_config == "" or torchao_config is None: + return model + elif "int8wo" in torchao_config: + quantize_(model, int8_weight_only(), filter_fn=filter_fn) + elif "int8dq" in torchao_config: + quantize_(model, int8_dynamic_activation_int8_weight(), filter_fn=filter_fn) + elif "int4wo" in torchao_config: + group_size = int(torchao_config.split("-")[-1]) + assert group_size in [ + 32, + 64, + 128, + 256, + ], f"int4wo groupsize needs to be one of [32, 64, 128, 256] but got {group_size}" + quantize_(model, int4_weight_only(group_size=group_size), filter_fn=filter_fn) + elif "gemlite" in torchao_config: + # gemlite--- or + # gemlite-- (packing_bitwidth defaults to 32) + from gemlite.core import GemLiteLinearTriton + from torchao.quantization import gemlite_uintx_weight_only + + _quant_args = torchao_config.split("-") + bit_width = int(_quant_args[-2]) + group_size = None if _quant_args[-1] == "None" else int(_quant_args[-1]) + + try: + packing_bitwidth = int(_quant_args[-3]) + except (ValueError, IndexError): + # if only 2 inputs found or conversion fails, use default value + packing_bitwidth = 32 + + quantize_( + model, gemlite_uintx_weight_only(group_size, bit_width, packing_bitwidth) + ) + + # try to load gemlite kernel config + GemLiteLinearTriton.load_config(get_gemlite_cache_path()) + + elif "fp8wo" in torchao_config: + # this requires newer hardware + # [rank0]: AssertionError: fp8e4nv data type is not supported on CUDA arch < 89 + quantize_(model, float8_weight_only(), filter_fn=filter_fn) + elif "fp8dq" in torchao_config: + granularity = torchao_config.split("-")[-1] + GRANULARITY_MAP = { + "per_row": PerRow(), + "per_tensor": PerTensor(), + } + assert ( + granularity in GRANULARITY_MAP + ), f"Supported granularity are: {GRANULARITY_MAP.keys()}, got {granularity}" + quantize_( + model, + float8_dynamic_activation_float8_weight( + granularity=GRANULARITY_MAP[granularity] + ), + filter_fn=filter_fn, + ) + else: + raise ValueError(f"Unexpected config: {torchao_config}") + + return model diff --git a/python/sglang/srt/layers/utils.py b/python/sglang/srt/layers/utils.py new file mode 100644 index 000000000..d79ccc663 --- /dev/null +++ b/python/sglang/srt/layers/utils.py @@ -0,0 +1,36 @@ +import logging +import re +from functools import lru_cache + +import torch + +logger = logging.getLogger(__name__) + + +def get_layer_id(weight_name): + # example weight name: model.layers.10.self_attn.qkv_proj.weight + match = re.search(r"layers\.(\d+)\.", weight_name) + if match: + return int(match.group(1)) + return None + + +class PPMissingLayer(torch.nn.Identity): + # Adapted from + # https://github.com/vllm-project/vllm/blob/18ed3132d2bfe1df9a74729457b69243955221e8/vllm/model_executor/models/utils.py#L468C1-L486C1 + """ + A placeholder layer for missing layers in a pipeline parallel model. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + self.return_tuple = kwargs.get("return_tuple", False) + + def forward(self, *args, **kwargs): + """ + Return the first arg from args or the first value from kwargs. + + Wraps the input in a tuple if `self.return_tuple` is True. + """ + input = args[0] if args else next(iter(kwargs.values())) + return (input,) if self.return_tuple else input diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py new file mode 100644 index 000000000..66abb7541 --- /dev/null +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -0,0 +1,571 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.3.post1/vllm/model_executor/layers/vocab_parallel_embedding.py + +import logging +from dataclasses import dataclass +from typing import List, Optional, Sequence, Tuple + +import torch +from torch.nn.parameter import Parameter, UninitializedParameter + +from sglang.srt.distributed import ( + divide, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + parallel_state, + tensor_model_parallel_all_reduce, +) +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.layers.amx_utils import PackWeightMethod +from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size +from sglang.srt.layers.parameter import BasevLLMParameter +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, + method_has_implemented_embedding, +) +from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod +from sglang.srt.utils import ( + cpu_has_amx_support, + get_compiler_backend, + is_cpu, + set_weight_attrs, +) + +DEFAULT_VOCAB_PADDING_SIZE = 64 + +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() + +logger = logging.getLogger(__name__) + + +def pad_vocab_size(vocab_size: int, pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int: + """Pad the vocab size to the given value.""" + return ((vocab_size + pad_to - 1) // pad_to) * pad_to + + +def vocab_range_from_per_partition_vocab_size( + per_partition_vocab_size: int, rank: int, offset: int = 0 +) -> Sequence[int]: + index_f = rank * per_partition_vocab_size + index_l = index_f + per_partition_vocab_size + return index_f + offset, index_l + offset + + +def vocab_range_from_global_vocab_size( + global_vocab_size: int, rank: int, world_size: int, offset: int = 0 +) -> Sequence[int]: + per_partition_vocab_size = divide(global_vocab_size, world_size) + return vocab_range_from_per_partition_vocab_size( + per_partition_vocab_size, rank, offset=offset + ) + + +@dataclass +class VocabParallelEmbeddingShardIndices: + """Indices for a shard of a vocab parallel embedding.""" + + padded_org_vocab_start_index: int + padded_org_vocab_end_index: int + padded_added_vocab_start_index: int + padded_added_vocab_end_index: int + + org_vocab_start_index: int + org_vocab_end_index: int + added_vocab_start_index: int + added_vocab_end_index: int + + @property + def num_org_elements(self) -> int: + return self.org_vocab_end_index - self.org_vocab_start_index + + @property + def num_added_elements(self) -> int: + return self.added_vocab_end_index - self.added_vocab_start_index + + @property + def num_org_elements_padded(self) -> int: + return self.padded_org_vocab_end_index - self.padded_org_vocab_start_index + + @property + def num_added_elements_padded(self) -> int: + return self.padded_added_vocab_end_index - self.padded_added_vocab_start_index + + @property + def num_org_vocab_padding(self) -> int: + return self.num_org_elements_padded - self.num_org_elements + + @property + def num_added_vocab_padding(self) -> int: + return self.num_added_elements_padded - self.num_added_elements + + @property + def num_elements_padded(self) -> int: + return self.num_org_elements_padded + self.num_added_elements_padded + + def __post_init__(self): + # sanity checks + assert self.padded_org_vocab_start_index <= self.padded_org_vocab_end_index + assert self.padded_added_vocab_start_index <= self.padded_added_vocab_end_index + + assert self.org_vocab_start_index <= self.org_vocab_end_index + assert self.added_vocab_start_index <= self.added_vocab_end_index + + assert self.org_vocab_start_index <= self.padded_org_vocab_start_index + assert self.added_vocab_start_index <= self.padded_added_vocab_start_index + assert self.org_vocab_end_index <= self.padded_org_vocab_end_index + assert self.added_vocab_end_index <= self.padded_added_vocab_end_index + + assert self.num_org_elements <= self.num_org_elements_padded + assert self.num_added_elements <= self.num_added_elements_padded + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def get_masked_input_and_mask( + input_: torch.Tensor, + org_vocab_start_index: int, + org_vocab_end_index: int, + num_org_vocab_padding: int, + added_vocab_start_index: int, + added_vocab_end_index: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + # torch.compile will fuse all of the pointwise ops below + # into a single kernel, making it very fast + org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ < org_vocab_end_index) + added_vocab_mask = (input_ >= added_vocab_start_index) & ( + input_ < added_vocab_end_index + ) + added_offset = ( + added_vocab_start_index + - (org_vocab_end_index - org_vocab_start_index) + - num_org_vocab_padding + ) + valid_offset = (org_vocab_start_index * org_vocab_mask) + ( + added_offset * added_vocab_mask + ) + vocab_mask = org_vocab_mask | added_vocab_mask + input_ = vocab_mask * (input_ - valid_offset) + return input_, ~vocab_mask + + +class VocabParallelEmbedding(torch.nn.Module): + """Embedding parallelized in the vocabulary dimension. + + Adapted from torch.nn.Embedding, note that we pad the vocabulary size to + make sure it is divisible by the number of model parallel GPUs. + + In order to support various loading methods, we ensure that LoRA-added + embeddings are always at the end of TP-sharded tensors. In other words, + we shard base embeddings and LoRA embeddings separately (both padded), + and place them in the same tensor. + In this example, we will have the original vocab size = 1010, + added vocab size = 16 and padding to 64. Therefore, the total + vocab size with padding will be 1088 (because we first pad 1010 to + 1024, add 16, and then pad to 1088). + Therefore, the tensor format looks like the following: + TP1, rank 0 (no sharding): + |< --------BASE-------- >|< -BASE PADDING-- >|< -----LORA------ >|< -LORA PADDING-- >| + corresponding token_id: | 0 | 1 | ... | 1009 | -1 | ... | -1 | 1010 | ... | 1015 | -1 | ... | -1 | + index: | 0 | 1 | ... | 1009 | 1010 | ... | 1023 | 1024 | ... | 1039 | 1040 | ... | 1087 | + + TP2, rank 0: + |< --------------------BASE--------------------- >|< -----LORA------ >|< -LORA PADDING- >| + corresponding token_id: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 1000 | ... | 1015 | -1 | ... | -1 | + index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 527 | 520 | ... | 543 | + TP2, rank 1: + |< -----------BASE----------- >|< -BASE PADDING- >|< -----------LORA PADDING----------- >| + corresponding token_id: | 512 | 513 | 514 | ... | 1009 | -1 | ... | -1 | -1 | ... | -1 | -1 | ... | -1 | + index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 519 | 520 | ... | 543 | + + Args: + num_embeddings: vocabulary size. + embedding_dim: size of hidden state. + params_dtype: type of the parameters. + org_num_embeddings: original vocabulary size (without LoRA). + padding_size: padding size for the vocabulary. + quant_config: quant config for the layer + prefix: full name of the layer in the state dict + """ # noqa: E501 + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + *, + params_dtype: Optional[torch.dtype] = None, + org_num_embeddings: Optional[int] = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + enable_tp: bool = True, + use_attn_tp_group: bool = False, + use_presharded_weights: bool = False, + ): + super().__init__() + self.quant_config = quant_config + + self.enable_tp = enable_tp + if self.enable_tp: + if use_attn_tp_group: + tp_rank = get_attention_tp_rank() + self.tp_size = get_attention_tp_size() + else: + tp_rank = get_tensor_model_parallel_rank() + self.tp_size = get_tensor_model_parallel_world_size() + else: + assert use_attn_tp_group is False + tp_rank = 0 + self.tp_size = 1 + + self.num_embeddings = num_embeddings + self.org_vocab_size = org_num_embeddings or num_embeddings + + # Support the case where the vocab size is not divisible by the TP size. + if ( + _is_cpu + and pad_vocab_size(self.org_vocab_size, padding_size) % self.tp_size != 0 + ): + padding_size *= self.tp_size + self.padding_size = padding_size + + num_added_embeddings = num_embeddings - self.org_vocab_size + self.use_presharded_weights = use_presharded_weights + if use_presharded_weights: + assert ( + num_added_embeddings == 0 + ), "Lora is not supported with presharded weights." + + self.org_vocab_size_padded = pad_vocab_size( + self.org_vocab_size, self.padding_size + ) + self.num_embeddings_padded = pad_vocab_size( + self.org_vocab_size_padded + num_added_embeddings, self.padding_size + ) + assert self.org_vocab_size_padded <= self.num_embeddings_padded + + self.shard_indices = self._get_indices( + self.num_embeddings_padded, + self.org_vocab_size_padded, + self.num_embeddings, + self.org_vocab_size, + tp_rank, + self.tp_size, + ) + self.embedding_dim = embedding_dim + + quant_method = None + if quant_config is not None: + quant_method = quant_config.get_quant_method(self, prefix=prefix) + if quant_method is None: + quant_method = UnquantizedEmbeddingMethod() + + # If we are making an embedding layer, then our quantization linear + # method must implement the embedding operation. If we are another + # layer type like ParallelLMHead, this is not important. + is_embedding_layer = type(self.__class__) is VocabParallelEmbedding + quant_method_implements_embedding = method_has_implemented_embedding( + type(quant_method) + ) + if is_embedding_layer and not quant_method_implements_embedding: + raise NotImplementedError( + f"The class {type(quant_method).__name__} must implement " + "the 'embedding' method, see UnquantizedEmbeddingMethod." + ) + + self.quant_method: QuantizeMethodBase = quant_method + + if params_dtype is None: + params_dtype = torch.get_default_dtype() + # Divide the weight matrix along the vocaburaly dimension. + self.num_added_embeddings = self.num_embeddings - self.org_vocab_size + self.num_embeddings_per_partition = divide( + self.num_embeddings_padded, self.tp_size + ) + assert ( + self.shard_indices.num_elements_padded == self.num_embeddings_per_partition + ) + self.num_org_embeddings_per_partition = ( + self.shard_indices.org_vocab_end_index + - self.shard_indices.org_vocab_start_index + ) + self.num_added_embeddings_per_partition = ( + self.shard_indices.added_vocab_end_index + - self.shard_indices.added_vocab_start_index + ) + + self.quant_method.create_weights( + self, + self.embedding_dim, + [self.num_embeddings_per_partition], + self.embedding_dim, + self.num_embeddings_padded, + params_dtype=params_dtype, + weight_loader=self.weight_loader, + ) + + @classmethod + def _get_indices( + cls, + vocab_size_padded: int, + org_vocab_size_padded: int, + vocab_size: int, + org_vocab_size: int, + tp_rank: int, + tp_size: int, + ) -> VocabParallelEmbeddingShardIndices: + """Get start and end indices for vocab parallel embedding, following the + layout outlined in the class docstring, based on the given tp_rank and + tp_size.""" + num_added_embeddings_padded = vocab_size_padded - org_vocab_size_padded + padded_org_vocab_start_index, padded_org_vocab_end_index = ( + vocab_range_from_global_vocab_size(org_vocab_size_padded, tp_rank, tp_size) + ) + padded_added_vocab_start_index, padded_added_vocab_end_index = ( + vocab_range_from_global_vocab_size( + num_added_embeddings_padded, tp_rank, tp_size, offset=org_vocab_size + ) + ) + # remove padding + org_vocab_start_index = min(padded_org_vocab_start_index, org_vocab_size) + org_vocab_end_index = min(padded_org_vocab_end_index, org_vocab_size) + added_vocab_start_index = min(padded_added_vocab_start_index, vocab_size) + added_vocab_end_index = min(padded_added_vocab_end_index, vocab_size) + return VocabParallelEmbeddingShardIndices( + padded_org_vocab_start_index, + padded_org_vocab_end_index, + padded_added_vocab_start_index, + padded_added_vocab_end_index, + org_vocab_start_index, + org_vocab_end_index, + added_vocab_start_index, + added_vocab_end_index, + ) + + def get_sharded_to_full_mapping(self) -> Optional[List[int]]: + """Get a mapping that can be used to reindex the gathered + logits for sampling. + + During sampling, we gather logits from all ranks. The relationship + of index->token_id will follow the same format as outlined in the class + docstring. However, after the gather, we want to reindex the final + logits tensor to map index->token_id one-to-one (the index is always + equal the token_id it corresponds to). The indices returned by this + method allow us to do that. + """ + if self.tp_size < 2: + return None + + base_embeddings: List[int] = [] + added_embeddings: List[int] = [] + padding: List[int] = [] + for tp_rank in range(self.tp_size): + shard_indices = self._get_indices( + self.num_embeddings_padded, + self.org_vocab_size_padded, + self.num_embeddings, + self.org_vocab_size, + tp_rank, + self.tp_size, + ) + range_start = self.num_embeddings_per_partition * tp_rank + range_end = self.num_embeddings_per_partition * (tp_rank + 1) + base_embeddings.extend( + range(range_start, range_start + shard_indices.num_org_elements) + ) + padding.extend( + range( + range_start + shard_indices.num_org_elements, + range_start + shard_indices.num_org_elements_padded, + ) + ) + added_embeddings.extend( + range( + range_start + shard_indices.num_org_elements_padded, + range_start + + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements, + ) + ) + padding.extend( + range( + range_start + + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements, + range_start + + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements_padded, + ) + ) + assert ( + range_start + + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements_padded + == range_end + ) + ret = base_embeddings + added_embeddings + padding + assert len(ret) == self.num_embeddings_padded + return ret + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + output_dim = getattr(param, "output_dim", None) + packed_dim = getattr(param, "packed_dim", None) + + # If the parameter is a gguf weight, then load it directly. + if getattr(param, "is_gguf_weight_type", None): + param.data.copy_(loaded_weight) + param.weight_type = loaded_weight.item() + return + elif isinstance(param, UninitializedParameter): + shape = list(loaded_weight.shape) + if output_dim is not None: + shape[output_dim] = shape[output_dim] // self.tp_size + param.materialize(tuple(shape), dtype=loaded_weight.dtype) + + # If parameter does not have output dim, then it should + # be copied onto all gpus (e.g. g_idx for act_order gptq). + if output_dim is None: + assert param.data.shape == loaded_weight.shape + param.data.copy_(loaded_weight) + return + + # Shard indexes for loading the weight + start_idx = self.shard_indices.org_vocab_start_index + shard_size = self.shard_indices.org_vocab_end_index - start_idx + + # If param packed on the same dim we are sharding on, then + # need to adjust offsets of loaded weight by pack_factor. + if packed_dim is not None and packed_dim == output_dim: + packed_factor = ( + param.packed_factor + if isinstance(param, BasevLLMParameter) + else param.packed_factor + ) + assert loaded_weight.shape[output_dim] == ( + self.org_vocab_size // param.packed_factor + ) + start_idx = start_idx // packed_factor + shard_size = shard_size // packed_factor + else: + assert loaded_weight.shape[output_dim] == ( + self.org_vocab_size + // (self.tp_size if self.use_presharded_weights else 1) + ), f"{self.org_vocab_size=} {self.use_presharded_weights=} {loaded_weight.shape[output_dim]=}" + + # Copy the data. + if not self.use_presharded_weights: + loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) + param[: loaded_weight.shape[0]].data.copy_(loaded_weight) + param[loaded_weight.shape[0] :].data.fill_(0) + + def forward(self, input_): + if self.tp_size > 1: + # Build the mask. + masked_input, input_mask = get_masked_input_and_mask( + input_, + self.shard_indices.org_vocab_start_index, + self.shard_indices.org_vocab_end_index, + self.shard_indices.num_org_vocab_padding, + self.shard_indices.added_vocab_start_index, + self.shard_indices.added_vocab_end_index, + ) + else: + masked_input = input_ + # Get the embeddings. + with use_symmetric_memory(parallel_state.get_tp_group()) as sm: + output_parallel = self.quant_method.embedding(self, masked_input.long()) + sm.tag(output_parallel) + # Mask the output embedding. + if self.tp_size > 1: + output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) + # Reduce across all the model parallel GPUs. + output = tensor_model_parallel_all_reduce(output_parallel) + else: + output = output_parallel + return output + + def extra_repr(self) -> str: + s = f"num_embeddings={self.num_embeddings_per_partition}" + s += f", embedding_dim={self.embedding_dim}" + s += f", org_vocab_size={self.org_vocab_size}" + s += f", num_embeddings_padded={self.num_embeddings_padded}" + if self.enable_tp: + s += f", tp_size={self.tp_size}" + return s + + +class ParallelLMHead(VocabParallelEmbedding): + """Parallelized LM head. + + Output logits weight matrices used in the Sampler. The weight and bias + tensors are padded to make sure they are divisible by the number of + model parallel GPUs. + + Args: + num_embeddings: vocabulary size. + embedding_dim: size of hidden state. + bias: whether to use bias. + params_dtype: type of the parameters. + org_num_embeddings: original vocabulary size (without LoRA). + padding_size: padding size for the vocabulary. + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + *, + bias: bool = False, + params_dtype: Optional[torch.dtype] = None, + org_num_embeddings: Optional[int] = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + use_attn_tp_group: bool = False, + use_presharded_weights: bool = False, + ): + super().__init__( + num_embeddings, + embedding_dim, + params_dtype=params_dtype, + org_num_embeddings=org_num_embeddings, + padding_size=padding_size, + quant_config=quant_config, + prefix=prefix, + use_attn_tp_group=use_attn_tp_group, + use_presharded_weights=use_presharded_weights, + ) + self.quant_config = quant_config + + # We only support pack LMHead if it's not quantized. + if _is_cpu and _is_cpu_amx_available: + if hasattr(self, "weight") and self.weight.dtype == torch.bfloat16: + self.quant_method = PackWeightMethod(weight_names=["weight"]) + + if bias: + self.bias = Parameter( + torch.empty(self.num_embeddings_per_partition, dtype=params_dtype) + ) + set_weight_attrs( + self.bias, + { + "output_dim": 0, + "weight_loader": self.weight_loader, + }, + ) + else: + self.register_parameter("bias", None) + + def tie_weights(self, embed_tokens: VocabParallelEmbedding): + """Tie the weights with word embeddings.""" + # GGUF quantized embed_tokens. + if self.quant_config and self.quant_config.get_name() == "gguf": + return embed_tokens + else: + self.weight = embed_tokens.weight + return self + + def forward(self, input_): + del input_ + raise RuntimeError("LMHead's weights should be used in the sampler.") diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py new file mode 100644 index 000000000..7c2c232d5 --- /dev/null +++ b/python/sglang/srt/lora/backend/base_backend.py @@ -0,0 +1,155 @@ +from typing import Optional, Tuple, Union + +import torch + +from sglang.srt.lora.utils import LoRABatchInfo +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +class BaseLoRABackend: + """Base class for different Lora backends. + Each backend has its own implementation of Lora kernels. + + Args: + max_loras_per_batch: maximum number of different lora weights + that can be applied in a single forward batch. + device: the device where the backend runs. + """ + + def __init__(self, max_loras_per_batch: int, device: torch.device): + self.max_loras_per_batch = max_loras_per_batch + self.device = device + + def run_lora_a_sgemm( + self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs + ) -> torch.Tensor: + """Run segment Gemm of lora a modules with current backend. + The definition of segment Gemm can be referred to https://docs.flashinfer.ai/api/gemm.html. + + Args: + x: input matrix with shape (s, input_dim), here s is the sum of all sequence lengths + weights: a set of lora weights with shape (num_lora, c * r, input_dim), + here r is lora rank, c is a multiplier for stacked modules (e.g., c=3 for qkv_proj, c=2 for gate_up_proj) + usually input_dim is much larger than r + Returns: + result with shape (s, c * r) + """ + pass + + def run_lora_b_sgemm( + self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs + ) -> torch.Tensor: + """Run segment Gemm of lora b modules with current backend. + The definition of segment Gemm can be referred to https://docs.flashinfer.ai/api/gemm.html. + + Args: + x: input matrix with shape (s, r), here s is the sum of all sequence lengths, r is lora rank + weights: a set of lora weights with shape (num_lora, output_dim, r) + usually output_dim is much larger than r + Returns: + result with shape (s, output_dim) + """ + pass + + def run_qkv_lora( + self, + x: torch.Tensor, + qkv_lora_a: torch.Tensor, + qkv_lora_b: Union[torch.Tensor, Tuple[torch.Tensor]], + *args, + **kwargs, + ) -> torch.Tensor: + """Run the lora pass for QKV Layer. + + Args: + x: input matrix with shape (s, input_dim), here s is the sum of all sequence lengths + qkv_lora_a: lora_a module for qkv, with shape (num_lora, 3 * r, input_dim) + qkv_lora_b: lora_b module for qkv. + If passed in as a tensor, its shape should be (num_lora,output_dim_q + 2 * output_dim_kv, r) + If passed in as a tuple of two tensors, it should contain: + a lora_b module for q, with shape (1, num_lora, output_dim_q, r) + and a combined lora_b module for kv, with shape (2, num_lora, output_dim_kv, r) + Returns: + result with shape (s, output_dim_q + 2 * output_dim_kv) + """ + pass + + def run_gate_up_lora( + self, + x: torch.Tensor, + gate_up_lora_a: torch.Tensor, + gate_up_lora_b: Union[torch.Tensor, Tuple[torch.Tensor]], + *args, + **kwargs, + ) -> torch.Tensor: + """Run the lora pass for gate_up_proj, usually attached to MergedColumnParallelLayer. + + Args: + x: input matrix with shape (s, input_dim), here s is the sum of all sequence lengths + gate_up_lora_a: lora_a module for gate_up_proj, with shape (num_lora, 2 * r, input_dim) + gate_up_lora_b: lora_b module for qkv. + If passed in as a tensor, its shape should be (num_lora, 2 * output_dim, r) + If passed in as a tuple, it should contain two tensors with shape (num_lora, output_dim, r) + Returns: + result with shape (s, 2 * output_dim) + """ + pass + + def init_cuda_graph_batch_info( + self, + cuda_graph_batch_info: LoRABatchInfo, + max_bs_in_cuda_graph: int, + ): + """Initialize the batch info for CUDA Graph mode. + + This method provides a hook for each backend to conduct its own initialization + logic for CUDA Graph mode. + + Args: + cuda_graph_batch_info: the LoRABatchInfo object created in LoraManager + max_bs_in_cuda_graph: maximum batch size for CUDA Graph mode + """ + pass + + def prepare_lora_batch( + self, + forward_batch: ForwardBatch, + weight_indices: list[int], + lora_ranks: list[int], + scalings: list[float], + batch_info: Optional[LoRABatchInfo] = None, + ): + """Prepare the lora weights and batch info for current forward batch. + + This method provides a hook for each backend to conduct its own preparation + logic for each forward batch. + + Args: + forward_batch: the ForwardBatch object for current forward pass + weight_indices: list of indices of lora weights to be applied for current batch + lora_ranks: list of lora ranks corresponding to weight_indices + scalings: list of scaling factors corresponding to weight_indices + batch_info: optional LoRABatchInfo object, if not provided, the backend should use its own + internal batch info (e.g., self.cuda_graph_batch_info for CUDA Graph mode) + """ + pass + + +def get_backend_from_name(name: str) -> BaseLoRABackend: + """ + Get corresponding backend class from backend's name + """ + if name == "triton": + from sglang.srt.lora.backend.triton_backend import TritonLoRABackend + + return TritonLoRABackend + # elif name == "csgmv": + # from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend + + # return ChunkedSgmvLoRABackend + elif name == "flashinfer": + raise ValueError( + "FlashInfer LoRA backend has been deprecated, please use `triton` instead." + ) + else: + raise ValueError(f"Invalid backend: {name}") diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py new file mode 100644 index 000000000..7abeef770 --- /dev/null +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -0,0 +1,176 @@ +from typing import Optional + +import torch + +from sglang.srt.lora.backend.base_backend import BaseLoRABackend +from sglang.srt.lora.triton_ops import ( + gate_up_lora_b_fwd, + qkv_lora_b_fwd, + sgemm_lora_a_fwd, + sgemm_lora_b_fwd, +) +from sglang.srt.lora.utils import LoRABatchInfo +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +class TritonLoRABackend(BaseLoRABackend): + name = "triton" + + def __init__(self, max_loras_per_batch: int, device: torch.device): + super().__init__(max_loras_per_batch, device) + + def run_lora_a_sgemm( + self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs + ) -> torch.Tensor: + return sgemm_lora_a_fwd(x, weights, self.batch_info) + + def run_lora_b_sgemm( + self, + x: torch.Tensor, + weights: torch.Tensor, + base_output: torch.Tensor = None, + *args, + **kwargs + ) -> torch.Tensor: + return sgemm_lora_b_fwd(x, weights, self.batch_info, base_output) + + def run_qkv_lora( + self, + x: torch.Tensor, + qkv_lora_a: torch.Tensor, + qkv_lora_b: torch.Tensor, + output_offset: torch.Tensor, + max_qkv_out_dim: int, + base_output: torch.Tensor = None, + *args, + **kwargs + ) -> torch.Tensor: + + # x: (s, input_dim) + # qkv_lora_a: (num_lora, 3 * r, input_dim) + # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r) + assert isinstance(qkv_lora_b, torch.Tensor) + + lora_a_output = sgemm_lora_a_fwd(x, qkv_lora_a, self.batch_info, stack_num=3) + lora_output = qkv_lora_b_fwd( + lora_a_output, + qkv_lora_b, + self.batch_info, + output_offset, + max_qkv_out_dim, + base_output, + ) + return lora_output + + def run_gate_up_lora( + self, + x: torch.Tensor, + gate_up_lora_a: torch.Tensor, + gate_up_lora_b: torch.Tensor, + base_output: torch.Tensor = None, + *args, + **kwargs + ) -> torch.Tensor: + + # x: (s, input_dim) + # gate_up_lora_a: (num_lora, 2 * r, input_dim) + # gate_up_lora_b: (num_lora, 2 * output_dim, r) + assert isinstance(gate_up_lora_b, torch.Tensor) + output_dim = gate_up_lora_b.shape[-2] // 2 + + # lora_a_output: (s, 2 * r) + lora_a_output = sgemm_lora_a_fwd( + x, gate_up_lora_a, self.batch_info, stack_num=2 + ) + lora_output = gate_up_lora_b_fwd( + lora_a_output, + gate_up_lora_b, + self.batch_info, + output_dim, + base_output, + ) + return lora_output + + def init_cuda_graph_batch_info( + self, cuda_graph_batch_info: LoRABatchInfo, max_bs_in_cuda_graph: int + ): + # Initialize seg_lens and seg_indptr for CUDA graph as they remain constant + # across batches. + cuda_graph_batch_info.seg_lens[:max_bs_in_cuda_graph].fill_(1) + torch.cumsum( + cuda_graph_batch_info.seg_lens[:max_bs_in_cuda_graph], + dim=0, + out=cuda_graph_batch_info.seg_indptr[1 : max_bs_in_cuda_graph + 1], + ) + + def prepare_lora_batch( + self, + forward_batch: ForwardBatch, + weight_indices: list[int], + lora_ranks: list[int], + scalings: list[float], + batch_info: Optional[LoRABatchInfo] = None, + ): + # Use pinned memory to avoid synchronizations during host-to-device transfer + weight_indices_tensor = torch.tensor( + weight_indices, dtype=torch.int32, pin_memory=True, device="cpu" + ) + lora_ranks_tensor = torch.tensor( + lora_ranks, dtype=torch.int32, pin_memory=True, device="cpu" + ) + scalings_tensor = torch.tensor( + scalings, dtype=torch.float, pin_memory=True, device="cpu" + ) + + bs = forward_batch.batch_size + + if batch_info is not None: + assert ( + batch_info.use_cuda_graph + ), "batch_info.use_cuda_graph must be True when batch_info is provided" + batch_info.bs = forward_batch.batch_size + batch_info.num_segments = forward_batch.batch_size + else: + max_len = ( + # Calculate max_len from the CPU copy to avoid D2H transfer. + max(forward_batch.extend_seq_lens_cpu) + if forward_batch.forward_mode.is_extend() + else 1 + ) + seg_lens = ( + forward_batch.extend_seq_lens + if forward_batch.forward_mode.is_extend() + else torch.ones(bs, device=self.device) + ) + seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device) + seg_indptr[1:] = torch.cumsum(seg_lens, dim=0) + + batch_info = LoRABatchInfo( + bs=forward_batch.batch_size, + num_segments=forward_batch.batch_size, + max_len=max_len, + use_cuda_graph=False, + seg_lens=seg_lens, + seg_indptr=seg_indptr, + weight_indices=torch.empty( + (bs,), dtype=torch.int32, device=self.device + ), + lora_ranks=torch.empty( + (self.max_loras_per_batch,), dtype=torch.int64, device=self.device + ), + scalings=torch.empty( + (self.max_loras_per_batch,), dtype=torch.float, device=self.device + ), + permutation=None, + ) + + # Copy to device asynchronously + batch_info.lora_ranks[: self.max_loras_per_batch].copy_( + lora_ranks_tensor, non_blocking=True + ) + batch_info.scalings[: self.max_loras_per_batch].copy_( + scalings_tensor, non_blocking=True + ) + batch_info.weight_indices[:bs].copy_(weight_indices_tensor, non_blocking=True) + + self.batch_info = batch_info diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py new file mode 100644 index 000000000..4426faccb --- /dev/null +++ b/python/sglang/srt/lora/layers.py @@ -0,0 +1,352 @@ +import torch +from torch import nn + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.lora.backend.base_backend import BaseLoRABackend + + +class BaseLayerWithLoRA(nn.Module): + def __init__( + self, + base_layer: nn.Module, + lora_backend: BaseLoRABackend, + ): + super().__init__() + self.base_layer: nn.Module = base_layer + self.set_lora: bool = False + self.lora_backend: BaseLoRABackend = lora_backend + + def forward(self, x: torch.Tensor): + return self.base_layer.forward(x) + + def set_lora_info(self, *args): + pass + + def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + pass + + def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + pass + + +class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): + """ + Vocab parallel embedding layer with support for LoRA (Low-Rank Adaptation). + + Note: The current version does not yet implement the LoRA functionality. + This class behaves exactly the same as the base VocabParallelEmbedding. + Future versions will integrate LoRA functionality to support efficient parameter fine-tuning. + """ + + def __init__( + self, + base_layer: VocabParallelEmbedding, + lora_backend: BaseLoRABackend, + ) -> None: + super().__init__(base_layer, lora_backend) + self.weight = base_layer.weight + + +class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA): + def __init__( + self, + base_layer: ColumnParallelLinear, + lora_backend: BaseLoRABackend, + ) -> None: + super().__init__(base_layer, lora_backend) + shard_size = self.base_layer.output_partition_sizes[0] + self.output_offset = torch.tensor( + [ + 0, + shard_size, + ], + dtype=torch.int32, + device=next(self.base_layer.parameters()).device, + ) + + def set_lora_info( + self, + A_buffer: torch.Tensor, + B_buffer: torch.Tensor, + ): + self.set_lora = True + self.A_buffer = A_buffer + self.B_buffer = B_buffer + + def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + lora_a_output = self.lora_backend.run_lora_a_sgemm(x, self.A_buffer) + lora_output = self.lora_backend.run_lora_b_sgemm( + x=lora_a_output, + weights=self.B_buffer, + output_offset=self.output_offset, + base_output=base_output, + ) + return lora_output + + def forward(self, input_: torch.Tensor): + # duplicate the logic in ColumnParallelLinear + bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None + output_parallel = self.base_layer.quant_method.apply( + self.base_layer, input_, bias + ) + + if self.set_lora: + output_parallel = self.apply_lora(output_parallel, input_) + + if self.base_layer.gather_output: + output = tensor_model_parallel_all_gather(output_parallel) + else: + output = output_parallel + output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None + return output, output_bias + + def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + return A + + def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + shard_size = self.base_layer.output_partition_sizes[0] + start_idx = tp_rank * shard_size + end_idx = (tp_rank + 1) * shard_size + B = B[start_idx:end_idx, :] + return B + + +class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): + def __init__( + self, + base_layer: MergedColumnParallelLinear, + lora_backend: BaseLoRABackend, + ) -> None: + super().__init__(base_layer, lora_backend) + + def set_lora_info( + self, + A_buffer: torch.Tensor, + B_buffer: torch.Tensor, + ): + self.set_lora = True + self.A_buffer_gate_up = A_buffer + self.B_buffer_gate_up = B_buffer + + shard_size = self.base_layer.output_partition_sizes[0] + self.output_offset = torch.tensor( + [ + 0, + shard_size, + 2 * shard_size, + ], + dtype=torch.int32, + device=next(self.base_layer.parameters()).device, + ) + + def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + lora_output = self.lora_backend.run_gate_up_lora( + x=x, + gate_up_lora_a=self.A_buffer_gate_up, + gate_up_lora_b=self.B_buffer_gate_up, + output_offset=self.output_offset, + base_output=base_output, + ) + return lora_output + + def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + return A + + def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + # Since the outputs for both gate and up are identical, we use a random one. + shard_size = self.base_layer.output_partition_sizes[0] + gate_size = self.base_layer.output_sizes[0] + start_idx = tp_rank * shard_size + end_idx = (tp_rank + 1) * shard_size + return torch.concat( + ( + B[start_idx:end_idx, :], + B[gate_size + start_idx : gate_size + end_idx], + ), + dim=0, + ) + + +class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): + def __init__( + self, + base_layer: QKVParallelLinear, + lora_backend: BaseLoRABackend, + ) -> None: + super().__init__(base_layer, lora_backend) + q_proj_shard_size = self.base_layer.q_proj_shard_size + kv_proj_shard_size = self.base_layer.kv_proj_shard_size + self.output_offset = torch.tensor( + [ + 0, + q_proj_shard_size, + q_proj_shard_size + kv_proj_shard_size, + q_proj_shard_size + 2 * kv_proj_shard_size, + ], + dtype=torch.int32, + device=next(self.base_layer.parameters()).device, + ) + + # For computing number of launched blocks + self.max_qkv_out_dim = max(q_proj_shard_size, kv_proj_shard_size) + + def set_lora_info( + self, + A_buffer_qkv: torch.Tensor, + B_buffer_qkv: torch.Tensor, + ): + self.set_lora = True + self.A_buffer_qkv = A_buffer_qkv + self.B_buffer_qkv = B_buffer_qkv + + def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + lora_output = self.lora_backend.run_qkv_lora( + x=x, + qkv_lora_a=self.A_buffer_qkv, + qkv_lora_b=self.B_buffer_qkv, + base_output=base_output, + output_offset=self.output_offset, + max_qkv_out_dim=self.max_qkv_out_dim, + ) + return lora_output + + def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + return A + + def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int) -> torch.Tensor: + base_layer = self.base_layer + q_proj_shard_size = base_layer.q_proj_shard_size + kv_proj_shard_size = base_layer.kv_proj_shard_size + num_kv_head_replicas = base_layer.num_kv_head_replicas + + q_start_idx = q_proj_shard_size * tp_rank + q_end_idx = q_start_idx + q_proj_shard_size + + kv_shard_id = tp_rank // num_kv_head_replicas + kv_start_idx = kv_proj_shard_size * kv_shard_id + kv_end_idx = kv_start_idx + kv_proj_shard_size + + q_size, k_size, _ = base_layer.output_sizes + B_q_shard = B[q_start_idx:q_end_idx, :] + B_k_shard = B[q_size + kv_start_idx : q_size + kv_end_idx, :] + B_v_shard = B[q_size + k_size + kv_start_idx : q_size + k_size + kv_end_idx, :] + + return torch.concat( + ( + B_q_shard, + B_k_shard, + B_v_shard, + ), + dim=0, + ) + + +class RowParallelLinearWithLoRA(BaseLayerWithLoRA): + def __init__( + self, + base_layer: RowParallelLinear, + lora_backend: BaseLoRABackend, + ) -> None: + super().__init__(base_layer, lora_backend) + + def set_lora_info(self, A_buffer: torch.Tensor, B_buffer: torch.Tensor): + self.set_lora = True + self.A_buffer = A_buffer + self.B_buffer = B_buffer + output_size = self.base_layer.output_size + self.output_offset = torch.tensor( + [ + 0, + output_size, + ], + dtype=torch.int32, + device=next(self.base_layer.parameters()).device, + ) + + def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + lora_a_output = self.lora_backend.run_lora_a_sgemm(x, self.A_buffer) + lora_output = self.lora_backend.run_lora_b_sgemm( + x=lora_a_output, + weights=self.B_buffer, + output_offset=self.output_offset, + base_output=base_output, + ) + return lora_output + + def forward(self, input_: torch.Tensor, skip_all_reduce=False): + # duplicate the logic in RowParallelLinear + if self.base_layer.input_is_parallel: + input_parallel = input_ + else: + tp_rank = get_tensor_model_parallel_rank() + splitted_input = split_tensor_along_last_dim( + input_, num_partitions=self.base_layer.tp_size + ) + input_parallel = splitted_input[tp_rank].contiguous() + output_parallel = self.base_layer.quant_method.apply( + self.base_layer, input_parallel + ) + + if self.set_lora: + output_parallel = self.apply_lora(output_parallel, input_parallel) + + if ( + self.base_layer.reduce_results + and self.base_layer.tp_size > 1 + and not skip_all_reduce + ): + output_ = tensor_model_parallel_all_reduce(output_parallel) + else: + output_ = output_parallel + + if not self.base_layer.skip_bias_add: + output = ( + output_ + self.base_layer.bias + if self.base_layer.bias is not None + else output_ + ) + output_bias = None + else: + output = output_ + output_bias = self.base_layer.bias + return output, output_bias + + def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int): + shard_size = self.base_layer.input_size_per_partition + start_idx = tp_rank * shard_size + end_idx = (tp_rank + 1) * shard_size + A = A[:, start_idx:end_idx].contiguous() + return A + + def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int): + return B + + +def get_lora_layer( + layer: nn.Module, lora_backend: BaseLoRABackend +) -> BaseLayerWithLoRA: + supported_layer_types = { + # the order matters + VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA, + QKVParallelLinear: QKVParallelLinearWithLoRA, + MergedColumnParallelLinear: MergedColumnParallelLinearWithLoRA, + ColumnParallelLinear: ColumnParallelLinearWithLoRA, + RowParallelLinear: RowParallelLinearWithLoRA, + } + for src_layer_type, lora_layer_type in supported_layer_types.items(): + if isinstance(layer, src_layer_type): # pylint: disable=unidiomatic-typecheck + ret = lora_layer_type(layer, lora_backend) + return ret + raise Exception(f"No corresponding LoRA layer supported for {type(layer)}.") diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py new file mode 100644 index 000000000..e7569624c --- /dev/null +++ b/python/sglang/srt/lora/lora.py @@ -0,0 +1,178 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== + +# Integrates "S-LoRA: Serving Thousands of Concurrent LoRA Adapters" +# and "Punica: Multi-Tenant LoRA Serving" + +# LoRA layers class inheritance adapted from: +# https://github.com/vllm-project/vllm/blob/4abf6336ec65c270343eb895e7b18786e9274176/vllm/lora/layers.py + +import logging +import re +from typing import Dict, List + +import torch +from torch import nn + +from sglang.srt.configs.load_config import LoadConfig +from sglang.srt.hf_transformers_utils import AutoConfig +from sglang.srt.lora.backend.base_backend import BaseLoRABackend + +# from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend +from sglang.srt.lora.backend.triton_backend import TritonLoRABackend +from sglang.srt.lora.lora_config import LoRAConfig +from sglang.srt.model_loader.loader import DefaultModelLoader + +logger = logging.getLogger(__name__) + + +class LoRALayer(nn.Module): + def __init__(self, config: LoRAConfig, base_hf_config: AutoConfig): + super().__init__() + self.config: LoRAConfig = config + self.base_hf_config: AutoConfig = base_hf_config + + # lora weights in cpu. The weights are loaded from checkpoint. + self.weights: Dict[str, torch.Tensor] = {} + + +class LoRAAdapter(nn.Module): + def __init__( + self, + uid: str, + config: LoRAConfig, + base_hf_config: AutoConfig, + load_config: LoadConfig, + lora_backend: BaseLoRABackend, + ): + super().__init__() + self.uid: str = uid + self.config: LoRAConfig = config + assert self.config.hf_config["peft_type"].lower() == "lora" + self.base_hf_config: AutoConfig = base_hf_config + self.load_config: LoadConfig = load_config + self.lora_backend: BaseLoRABackend = lora_backend + self.scaling: float = self.config.lora_alpha / self.config.r + + self.layers: List[LoRALayer] = nn.ModuleList( + [ + LoRALayer(config, base_hf_config) + for _ in range(base_hf_config.num_hidden_layers) + ] + ) + + self.weights: Dict[str, torch.Tensor] = {} + + # initialize the LoRA weights to cpu + def initialize_weights(self): + model_path = self.config.path + loader = DefaultModelLoader(self.load_config) + revision = getattr(self.config.hf_config, "revision", None) + for name, loaded_weight in loader._get_weights_iterator( + DefaultModelLoader.Source( + model_path, revision=revision, fall_back_to_pt=True + ) + ): + match = re.search(r"layers\.(\d+)\.", name) + if match is not None: + layer_id = int(match.group(1)) + self.layers[layer_id].weights[name] = loaded_weight.cpu() + else: + self.weights[name] = loaded_weight.cpu() + + # normalize kv_proj and gate_up_proj + for layer in self.layers: + weight_names = list(layer.weights.keys()) + self.normalize_qkv_proj(weight_names, layer.weights) + self.normalize_gate_up_proj(weight_names, layer.weights) + + def normalize_qkv_proj( + self, weight_names: List[str], weights: Dict[str, torch.Tensor] + ): + # Collect target q/k/v modules. This process is necessary since there might be no lora attached to k_proj + target_module = set() + for weight_name in weight_names: + if "k_proj" in weight_name: + target_module.add("k_proj") + if "q_proj" in weight_name: + target_module.add("q_proj") + if "v_proj" in weight_name: + target_module.add("v_proj") + if "qkv_proj" in weight_name: + target_module.add("qkv_proj") + if len(target_module) == 0: + return + + for weight_name in weight_names: + # We assume every lora adaptor should contain lora modules for q_proj + if "q_proj" in weight_name: + q_name = weight_name + k_name = weight_name.replace("q_proj", "k_proj") + v_name = weight_name.replace("q_proj", "v_proj") + qkv_name = weight_name.replace("q_proj", "qkv_proj") + + # If k_proj doesn't have lora, initialize it to zero + k_proj_weight = ( + weights[k_name] + if "k_proj" in target_module + else torch.zeros_like(weights[v_name]) + ) + weights[qkv_name] = torch.cat( + ( + weights[q_name], + k_proj_weight, + weights[v_name], + ), + 0, + ) + weights.pop(q_name) + if "k_proj" in target_module: + weights.pop(k_name) + weights.pop(v_name) + elif "qkv_proj" in weight_name: + # If qkv_proj is already stacked, we normalize it following the SGL convention. + qkv_name = weight_name + q_name = weight_name.replace("qkv_proj", "q_proj") + k_name = weight_name.replace("qkv_proj", "k_proj") + v_name = weight_name.replace("qkv_proj", "v_proj") + if "lora_A" in weight_name: + weights[qkv_name] = weights[qkv_name].repeat(3, 1) + # else: no-op as LoRA B weight is already stacked. + + def normalize_gate_up_proj( + self, weight_names: List[str], weights: Dict[str, torch.Tensor] + ): + for weight_name in weight_names: + if "gate_proj" in weight_name: + up_name = weight_name.replace("gate_proj", "up_proj") + gate_up_name = weight_name.replace("gate_proj", "gate_up_proj") + if up_name not in weights: + weights[up_name] = torch.zeros_like(weights[weight_name]) + assert isinstance(self.lora_backend, TritonLoRABackend), ( + f"LoRA weight initialization currently only supported for 'triton' backend. " + f"Received backend: {self.lora_backend.name}. Please verify your backend configuration " + f"or consider implementing custom initialization logic for other backends." + ) + weights[gate_up_name] = torch.cat( + (weights[weight_name], weights[up_name]), 0 + ) + weights.pop(weight_name) + if up_name in weights: + weights.pop(up_name) + elif "gate_up_proj" in weight_name: + # If gate_up_proj is already stacked, we normalize it following the SGL convention + gate_up_name = weight_name + if "lora_A" in weight_name: + weights[gate_up_name] = weights[gate_up_name].repeat(2, 1) + # else: no-op as LoRA B weight is already stacked. diff --git a/python/sglang/srt/lora/lora_config.py b/python/sglang/srt/lora/lora_config.py new file mode 100644 index 000000000..185b7b824 --- /dev/null +++ b/python/sglang/srt/lora/lora_config.py @@ -0,0 +1,47 @@ +# Copyright 2023-2024 SGLang 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 json +import os + +from huggingface_hub import snapshot_download + + +class LoRAConfig: + def __init__( + self, + path: str, + ) -> None: + self.path = path + self.hf_config = self.get_lora_config() + self.target_modules = self.hf_config["target_modules"] + + # TODO: Support more modules + if any(module in self.target_modules for module in ["embed_tokens", "lm_head"]): + raise ValueError("Not supported yet") + + self.r = self.hf_config["r"] + self.lora_alpha = self.hf_config["lora_alpha"] + + def get_lora_config(self, dummy=False): + if dummy: + raise NotImplementedError() + else: + if not os.path.isdir(self.path): + weights_dir = snapshot_download(self.path, allow_patterns=["*.json"]) + else: + weights_dir = self.path + config_name = "adapter_config.json" + with open(os.path.join(weights_dir, config_name), "r") as f: + return json.load(f) diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py new file mode 100644 index 000000000..baf120ca2 --- /dev/null +++ b/python/sglang/srt/lora/lora_manager.py @@ -0,0 +1,440 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== + +# Integrates "S-LoRA: Serving Thousands of Concurrent LoRA Adapters" +# and "Punica: Multi-Tenant LoRA Serving" + +import logging +from typing import Dict, Iterable, List, Optional, Set, Tuple + +import torch + +from sglang.srt.configs.load_config import LoadConfig +from sglang.srt.hf_transformers_utils import AutoConfig +from sglang.srt.lora.backend.base_backend import BaseLoRABackend, get_backend_from_name +from sglang.srt.lora.layers import BaseLayerWithLoRA, get_lora_layer +from sglang.srt.lora.lora import LoRAAdapter +from sglang.srt.lora.lora_config import LoRAConfig +from sglang.srt.lora.lora_registry import LoRARef +from sglang.srt.lora.mem_pool import LoRAMemoryPool +from sglang.srt.lora.utils import ( + LoRABatchInfo, + LoRAType, + get_layer_id, + get_normalized_target_modules, + get_target_module_name, +) +from sglang.srt.managers.io_struct import LoRAUpdateResult +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.utils import replace_submodule + +logger = logging.getLogger(__name__) + + +class LoRAManager: + def __init__( + self, + base_model: torch.nn.Module, + base_hf_config: AutoConfig, + max_loras_per_batch: int, + load_config: LoadConfig, + dtype: torch.dtype, + lora_backend: str = "triton", + tp_size: int = 1, + tp_rank: int = 0, + max_lora_rank: Optional[int] = None, + target_modules: Optional[Iterable[str]] = None, + lora_paths: Optional[List[LoRARef]] = None, + ): + self.base_model: torch.nn.Module = base_model + self.base_hf_config: AutoConfig = base_hf_config + self.max_loras_per_batch: int = max_loras_per_batch + self.load_config: LoadConfig = load_config + self.dtype: torch.dtype = dtype + self.device: torch.device = next(self.base_model.parameters()).device + self.tp_size: int = tp_size + self.tp_rank: int = tp_rank + + # LoRA backend for running sgemm kernels + logger.info(f"Using {lora_backend} as backend of LoRA kernels.") + backend_type = get_backend_from_name(lora_backend) + self.lora_backend: BaseLoRABackend = backend_type( + max_loras_per_batch=max_loras_per_batch, + device=self.device, + ) + + # Initialize mutable internal state of the LoRAManager. + self.init_state( + max_lora_rank=max_lora_rank, + target_modules=target_modules, + lora_paths=lora_paths, + ) + + def init_cuda_graph_batch_info(self, max_bs_in_cuda_graph: int): + self.max_bs_in_cuda_graph = max_bs_in_cuda_graph + with torch.device("cuda"): + self.cuda_graph_batch_info = LoRABatchInfo( + bs=max_bs_in_cuda_graph, + use_cuda_graph=True, + num_segments=None, + seg_lens=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32), + seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32), + max_len=1, + weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32), + permutation=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32), + lora_ranks=torch.zeros(self.max_loras_per_batch, dtype=torch.int32), + scalings=torch.zeros(self.max_loras_per_batch, dtype=torch.float), + ) + + self.lora_backend.init_cuda_graph_batch_info( + cuda_graph_batch_info=self.cuda_graph_batch_info, + max_bs_in_cuda_graph=max_bs_in_cuda_graph, + ) + + def create_lora_update_result( + self, success: bool, error_message: str = "" + ) -> LoRAUpdateResult: + return LoRAUpdateResult( + success=success, + error_message=error_message, + loaded_adapters={ + lora_ref.lora_name: lora_ref.lora_path + for lora_ref in self.lora_refs.values() + }, + ) + + def load_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateResult: + """ + Load a single LoRA adapter from the specified path. + + Args: + lora_ref (LoRARef): The LoRARef object containing the LoRA name, path, and ID. + """ + assert ( + lora_ref.lora_name is not None and lora_ref.lora_path is not None + ), "LoRARef must have both lora_name and lora_path set for loading." + assert ( + lora_ref.lora_id not in self.loras + ), f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend." + + try: + # load configs + new_adapter = LoRAConfig(lora_ref.lora_path) + self.validate_new_adapter(new_adapter, lora_ref) + self.configs[lora_ref.lora_id] = new_adapter + + # load weights + self.load_lora_weights(lora_ref) + + # keep metadata for displayed messages + self.lora_refs[lora_ref.lora_id] = lora_ref + self.num_pinned_loras += int(lora_ref.pinned) + except Exception as e: + return self.create_lora_update_result( + success=False, + error_message=str(e), + ) + + return self.create_lora_update_result(success=True) + + def validate_new_adapter(self, lora_config: LoRAConfig, lora_ref: LoRARef): + """ + Validate if an adapter can be loaded into the current LoRA memory pool and generate error if it is incompatible. + """ + + # Check if the LoRA adapter shape is compatible with the current LoRA memory pool configuration. + memory_pool = getattr(self, "memory_pool", None) + incompatible = memory_pool and not memory_pool.can_support(lora_config) + if incompatible: + raise ValueError( + f"LoRA adapter {lora_ref.lora_name} with rank {lora_config.r} is incompatible with the current " + "LoRA memory pool configuration. Please ensure that the LoRA adapter's rank is within the configured " + "`--max-lora-rank` and that the target modules are included in `--lora-target-modules`." + ) + + # Ensure pinned LoRA adapters does not exceed maximal limit or cause starvation. + if lora_ref.pinned and self.num_pinned_loras >= self.max_loras_per_batch - 1: + raise ValueError( + f"Failed to load LoRA adapter {lora_ref.lora_name} as a pinned adapter. It is not allowed to pin all slots " + "in the LoRA memory pool to avoid starvation for unpinned adapters and base models. Please increase your " + "`--max-loras-per-batch` or load it as unpinned LoRA adapters." + ) + + def unload_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateResult: + """ + Unload LoRA adapters by their names. This will remove the adapters from the memory pool and + delete the corresponding LoRA modules. + """ + + adapter = self.configs.get(lora_ref.lora_id) + lora_ref = self.lora_refs.get(lora_ref.lora_id) + assert ( + adapter is not None and lora_ref is not None + ), f"LoRA adapter with ID {lora_ref.lora_id} is not loaded. This should have been verified before request is sent to the backend." + + try: + del self.configs[lora_ref.lora_id] + del self.loras[lora_ref.lora_id] + del self.lora_refs[lora_ref.lora_id] + self.num_pinned_loras -= int(lora_ref.pinned) + except Exception as e: + return self.create_lora_update_result( + success=False, + error_message=str(e), + ) + + return self.create_lora_update_result(success=True) + + def validate_lora_batch(self, lora_ids: set[str]) -> bool: + """ + Validate if the LoRA IDs in the batch can be loaded into the current LoRA memory pool. + """ + if len(lora_ids) > self.max_loras_per_batch: + return False + + # skip pinned LoRA check if no pinned LoRA adapters are loaded. + if self.num_pinned_loras == 0: + return True + + # counting the number of pinned LoRA adapters in the batch. + pinned_loras_in_batch = 0 + for lora_id in lora_ids: + if lora_id is not None: + lora_ref = self.lora_refs.get(lora_id) + assert ( + lora_ref is not None + ), f"LoRA ID {lora_id} not found in lora_refs." + pinned_loras_in_batch += int(lora_ref.pinned) + + assert pinned_loras_in_batch <= self.num_pinned_loras, ( + f"Number of pinned LoRA adapters in the batch ({pinned_loras_in_batch}) exceeds the total number of pinned adapters " + f"({self.num_pinned_loras}). This indicates a bug in the LoRA loading logic." + ) + + required_slots = len(lora_ids) - pinned_loras_in_batch + mem_pool_vacancy = self.memory_pool.max_loras_per_batch - self.num_pinned_loras + + return required_slots <= mem_pool_vacancy + + def prepare_lora_batch(self, forward_batch: ForwardBatch): + # Load active loras into lora memory pool + cur_uids = set(forward_batch.lora_ids) + + assert len(cur_uids) <= self.max_loras_per_batch + self.memory_pool.prepare_lora_batch( + cur_uids=cur_uids, + lora_adapters=self.loras, + lora_modules=self.lora_modules, + lora_refs=self.lora_refs.copy(), # copy snapshot of current lora_refs to avoid mutation during the batch preparation. + ) + + # set up batch info shared by all lora modules + bs = forward_batch.batch_size + + use_cuda_graph = ( + hasattr(self, "max_bs_in_cuda_graph") + and bs <= self.max_bs_in_cuda_graph + and forward_batch.forward_mode.is_cuda_graph() + ) + + weight_indices = [0] * len(forward_batch.lora_ids) + lora_ranks = [0] * self.max_loras_per_batch + scalings = [0] * self.max_loras_per_batch + for i, uid in enumerate(forward_batch.lora_ids): + weight_indices[i] = self.memory_pool.get_buffer_id(uid) + if uid is not None: + lora = self.loras[uid] + lora_ranks[weight_indices[i]] = lora.config.r + scalings[weight_indices[i]] = lora.scaling + # Do in-place updates when CUDA graph is enabled and the batch forward mode + # could use CUDA graph. + self.lora_backend.prepare_lora_batch( + forward_batch=forward_batch, + weight_indices=weight_indices, + lora_ranks=lora_ranks, + scalings=scalings, + batch_info=self.cuda_graph_batch_info if use_cuda_graph else None, + ) + + def update_lora_info(self): + """ + Update all LoRA modules to associate them with the latest memory buffer. + """ + for layer_id, layer_modules in enumerate(self.lora_modules): + for module_name, module in layer_modules.items(): + target_module = get_target_module_name( + module_name, self.memory_pool.target_modules + ) + module.set_lora_info( + self.memory_pool.get_tensor( + target_module=target_module, + layer_id=layer_id, + lora_type=LoRAType.LORA_A, + ), + self.memory_pool.get_tensor( + target_module=target_module, + layer_id=layer_id, + lora_type=LoRAType.LORA_B, + ), + ) + + def init_state( + self, + max_lora_rank: Optional[int] = None, + target_modules: Optional[Iterable[str]] = None, + lora_paths: Optional[List[LoRARef]] = None, + ): + """ + Initialize the internal (mutable) state of the LoRAManager. + + When `lora_paths` is provided and not empty, it might be used for inferring LoRA shape info such as + the target modules and max_lora_rank. + """ + + assert lora_paths or ( + max_lora_rank is not None and target_modules is not None + ), "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." + + self.init_lora_adapters(lora_paths) + self.init_lora_shapes( + max_lora_rank=max_lora_rank, + target_modules=target_modules, + ) + self.init_lora_modules() + self.init_memory_pool() + self.update_lora_info() + + def init_lora_adapters(self, lora_paths: Optional[List[LoRARef]] = None): + # Configs of all active LoRA adapters, indexed by LoRA ID. + self.configs: Dict[str, LoRAConfig] = {} + + # LoRA adapter weights cached in CPU memory, indexed by LoRA ID. + self.loras: Dict[str, LoRAAdapter] = {} + + # Mapping from LoRA ID to LoRARef object. + self.lora_refs: Dict[str, LoRARef] = {} + + # Count of pinned LoRA adapters. + self.num_pinned_loras: int = 0 + + if lora_paths: + for lora_ref in lora_paths: + result = self.load_lora_adapter(lora_ref) + if not result.success: + raise RuntimeError( + f"Failed to load LoRA adapter {lora_ref.lora_name}: {result.error_message}" + ) + + def init_lora_shapes( + self, + max_lora_rank: Optional[int] = None, + target_modules: Optional[Iterable[str]] = None, + ): + """Infer LoRA target modules and max_lora_rank from loaded adapters if not provided.""" + + self.target_modules = ( + get_normalized_target_modules(target_modules) if target_modules else set() + ) + + for lora_id, config in self.configs.items(): + if not isinstance(config.target_modules, list): + raise ValueError( + f"SGLang currently only supports inferring LoRA target modules when a list of " + "suffixes is provided in `target_modules` field of PEFT config. Please explicitly " + "specify `--lora-target-modules` during server startup. You can specify `all` to " + "enable all support modules types. " + ) + + adapter_target_modules = get_normalized_target_modules( + config.target_modules + ) + + if target_modules is not None: + # When `--lora-target-modules` is provided, validate adapter target modules is a subset of the specified target modules. + if not adapter_target_modules.issubset(self.target_modules): + unsupported_modules = adapter_target_modules - self.target_modules + lora_name = self.lora_refs[lora_id].lora_name + raise ValueError( + f"LoRA adapter '{lora_name}' contains target modules {sorted(unsupported_modules)} " + f"that are not included in the specified --lora-target-modules {sorted(self.target_modules)}. " + f"Please update --lora-target-modules to include all required modules: " + f"{sorted(self.target_modules | adapter_target_modules)}, or use 'all' to enable all supported modules." + ) + else: + # Otherwise, infer target_modules from adapter configs. + self.target_modules.update(adapter_target_modules) + + if max_lora_rank is not None: + self.max_lora_rank = max_lora_rank + else: + self.max_lora_rank = max( + [x.r for x in self.configs.values()], + default=0, + ) + + def load_lora_weights(self, lora_ref: LoRARef): + """ + Load the weights of a LoRA adapter to CPU memory and conducts post-loading validation. + """ + lora_adapter = LoRAAdapter( + lora_ref.lora_id, + self.configs[lora_ref.lora_id], + self.base_hf_config, + self.load_config, + self.lora_backend, + ) + lora_adapter.initialize_weights() + self.loras[lora_ref.lora_id] = lora_adapter + + def init_memory_pool(self): + """(Re)initialize the LoRA memory pool based on the current configurations.""" + self.memory_pool = LoRAMemoryPool( + base_hf_config=self.base_hf_config, + max_loras_per_batch=self.max_loras_per_batch, + dtype=self.dtype, + tp_size=self.tp_size, + tp_rank=self.tp_rank, + max_lora_rank=self.max_lora_rank, + target_modules=self.target_modules, + base_model=self.base_model, + ) + + def set_lora_module(self, module_name, module): + lora_module = get_lora_layer(module, self.lora_backend) + replace_submodule(self.base_model, module_name, lora_module) + return lora_module + + def init_lora_modules(self): + # Look-up table that essentially maps (layer_index, module_name) to the corresponding LoRA module. + self.lora_modules: List[Dict[str, BaseLayerWithLoRA]] = [ + {} for _ in range(self.base_hf_config.num_hidden_layers) + ] + + for module_name, module in self.base_model.named_modules(): + # TODO (lifuhuang): in the future, we should consider generalizing the + # should_apply_lora function to support mapping by full module name instead + # of just the last part (e.g., "qkv_proj") to support scenarios with multiple + # attention stacks (e.g., multimodal models). + # See: https://github.com/sgl-project/sglang/issues/6608 + if getattr( + self.base_model, "should_apply_lora", None + ) and not self.base_model.should_apply_lora(module_name): + continue + + # The module should be converted if it is included in target_names + if module_name.split(".")[-1] in self.target_modules: + layer_id = get_layer_id(module_name) + self.lora_modules[layer_id][module_name] = self.set_lora_module( + module_name, module + ) diff --git a/python/sglang/srt/lora/lora_registry.py b/python/sglang/srt/lora/lora_registry.py new file mode 100644 index 000000000..51d2b0e66 --- /dev/null +++ b/python/sglang/srt/lora/lora_registry.py @@ -0,0 +1,207 @@ +# Copyright 2023-2024 SGLang 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 asyncio +from dataclasses import dataclass, field, fields +from typing import Dict, List, Optional, Union +from uuid import uuid4 + +from sglang.srt.aio_rwlock import RWLock +from sglang.srt.utils import ConcurrentCounter + + +@dataclass(frozen=True) +class LoRARef: + """ + Reference record for a LoRA model. + + This object guarantees a unique ``lora_id`` and may include ``lora_name``, ``lora_path``, and ``pinned``. + The ID eliminates conflicts from reused LoRA names or paths and can be used to generate deterministic cache + keys (e.g., radix cache). + """ + + lora_id: str = field(default_factory=lambda: uuid4().hex) + lora_name: Optional[str] = None + lora_path: Optional[str] = None + pinned: Optional[bool] = None + + def __post_init__(self): + if self.lora_id is None: + raise ValueError("lora_id cannot be None") + + def __str__(self) -> str: + parts = [ + f"{f.name}={value}" + for f in fields(self) + if (value := getattr(self, f.name)) is not None + ] + return f"{self.__class__.__name__}({', '.join(parts)})" + + +class LoRARegistry: + """ + The central registry to keep track of available LoRA adapters and ongoing LoRA requests. + + The `LoRARegistry` resides in the tokenizer manager process and acts as the single source of truth for all + available LoRA adapters. It supports concurrent inference and dynamic adapter updates through a two-phase + update / eventual consistency model between the tokenizer manager process and the scheduler processes. + """ + + def __init__(self, lora_paths: Optional[List[LoRARef]] = None): + assert lora_paths is None or all( + isinstance(lora, LoRARef) for lora in lora_paths + ), ( + "server_args.lora_paths should have been normalized to LoRARef objects during server initialization. " + "Please file an issue if you see this error." + ) + + # A read-write lock to ensure adapters loading / unloading operations are exclusive. + # Please note that the counter increment/decrement operations are not synchronized through this + # lock, as they are designed to be non-blocking and can be performed concurrently. + self._registry_lock = RWLock() + # A dictionary to hold LoRARef objects, mapping from LoRA name to LoRARef. + self._registry: Dict[str, LoRARef] = {} + # Counters for ongoing requests, mapping from LoRA ID to ConcurrentCounter. + self._counters: Dict[str, ConcurrentCounter] = {} + + # Initialize the registry with provided LoRA paths, if present. + if lora_paths: + for lora_ref in lora_paths: + self._register_adapter(lora_ref) + + async def register(self, lora_ref: LoRARef): + """ + Register a new LoRARef object in the registry. + + Args: + lora_ref (LoRARef): The LoRARef object to register. + """ + async with self._registry_lock.writer_lock: + self._register_adapter(lora_ref) + + async def unregister(self, lora_name: str) -> str: + """ + Unregister a LoRARef object from the registry and returns the removed LoRA ID. + + Args: + lora_name (str): The name of the LoRA model to unregister. + """ + async with self._registry_lock.writer_lock: + lora_ref = self._registry.get(lora_name, None) + if lora_ref is None: + raise ValueError( + f"LoRA with name {lora_name} does not exist. Loaded LoRAs: {self._registry.keys()}" + ) + del self._registry[lora_name] + + return lora_ref.lora_id + + async def acquire(self, lora_name: Union[str, List[str]]) -> Union[str, List[str]]: + """ + Queries registry for LoRA IDs based on LoRA names and start tracking the usage of the corresponding LoRA adapters + by incrementing its counter. + """ + + def _lookup(name: str) -> str: + if name is None: + return None + + lora_ref = self._registry.get(name, None) + if lora_ref is None: + raise ValueError( + f"The following requested LoRA adapters are not loaded: {name}\n" + f"Loaded adapters: {self._registry.keys()}." + ) + return lora_ref.lora_id + + async with self._registry_lock.reader_lock: + if isinstance(lora_name, str): + lora_id = _lookup(lora_name) + await self._counters[lora_id].increment(notify_all=False) + return lora_id + elif isinstance(lora_name, list): + lora_ids = [_lookup(name) for name in lora_name] + + # Increment the counters only after all IDs are looked up. + await asyncio.gather( + *[ + self._counters[id].increment(notify_all=False) + for id in lora_ids + if id is not None + ] + ) + return lora_ids + else: + raise TypeError( + "lora_name must be either a string or a list of strings." + ) + + async def release(self, lora_id: Union[str, List[str]]): + """ + Decrements the usage counter for a LoRA adapter, indicating that it is no longer in use. + """ + + async with self._registry_lock.reader_lock: + if isinstance(lora_id, str): + await self._counters[lora_id].decrement() + elif isinstance(lora_id, list): + await asyncio.gather( + *[ + self._counters[id].decrement() + for id in lora_id + if id is not None + ] + ) + else: + raise TypeError("lora_id must be either a string or a list of strings.") + + async def wait_for_unload(self, lora_id: str): + """ + Waits until the usage counter for a LoRA adapter reaches zero, indicating that it is no longer in use. + This is useful for ensuring that a LoRA adapter can be safely unloaded. + + This method itself is not synchronized, which is safe because it should only be called during LoRA unloading, + which itself is guaranteed to be sequential. + """ + assert ( + lora_id not in self._registry + ), "wait_for_unload should only be called after the LoRA adapter has been unregistered. " + assert ( + lora_id in self._counters + ), "The LoRA ID should still have a counter if it has been registered before." + + # Wait until no requests are using this LoRA adapter. + await self._counters[lora_id].wait_for_zero() + del self._counters[lora_id] + + def _register_adapter(self, lora_ref: LoRARef): + """ + Internal helper method to register a LoRA adapter. + """ + + if lora_ref.lora_name in self._registry: + raise ValueError( + f"LoRA with name {lora_ref.lora_name} already exists. Loaded LoRAs: {self._registry.keys()}" + ) + self._registry[lora_ref.lora_name] = lora_ref + self._counters[lora_ref.lora_id] = ConcurrentCounter() + return lora_ref + + @property + def num_registered_loras(self) -> int: + """ + Returns the total number of LoRA adapters currently registered. + """ + return len(self._registry) diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py new file mode 100644 index 000000000..cdf1707e8 --- /dev/null +++ b/python/sglang/srt/lora/mem_pool.py @@ -0,0 +1,309 @@ +import logging +from typing import Callable, Dict, Iterable, List, Optional, Set, Tuple, Union + +import torch + +from sglang.srt.distributed import divide +from sglang.srt.hf_transformers_utils import AutoConfig +from sglang.srt.lora.layers import BaseLayerWithLoRA +from sglang.srt.lora.lora import LoRAAdapter +from sglang.srt.lora.lora_config import LoRAConfig +from sglang.srt.lora.lora_registry import LoRARef +from sglang.srt.lora.utils import ( + ROW_PARALLELISM_LINEAR_LORA_NAMES, + LoRAType, + get_hidden_dim, + get_normalized_target_modules, + get_stacked_multiply, + get_target_module_name, +) + +logger = logging.getLogger(__name__) + + +class EmptySlot: + """ + Singleton class to represent an empty slot in the memory pool. + This is used to improve readability by not using special str as a placeholder. + """ + + __slots__ = () + + def __repr__(self): + return "|EMPTY|" + + def __new__(cls): + if not hasattr(cls, "_instance"): + cls._instance = super().__new__(cls) + return cls._instance + + +EMPTY_SLOT = EmptySlot() + + +class LoRAMemoryPool: + """Class for memory pool management of lora modules""" + + def __init__( + self, + base_hf_config: AutoConfig, + max_loras_per_batch: int, + dtype: torch.dtype, + tp_size: int, + tp_rank: int, + max_lora_rank: int, + target_modules: Set[str], + base_model: torch.nn.Module, + ): + self.base_hf_config: AutoConfig = base_hf_config + self.num_layer: int = base_hf_config.num_hidden_layers + self.max_loras_per_batch: int = max_loras_per_batch + self.dtype: torch.dtype = dtype + self.tp_size: int = tp_size + self.tp_rank: int = tp_rank + self.max_lora_rank: int = max_lora_rank + self.target_modules: Set[str] = target_modules + + # Both A_buffer and B_buffer maps lora weight names to its buffer space. + # A_buffer contains num_layer number of row-major tensors with shape + # (max_loras_per_batch, stacked_num * max_lora_dim, input_dim) + # B_buffer contains num_layer number of column-major tensors with shape + # (stacked_num, max_loras_per_batch, output_dim, max_lora_dim) + self.A_buffer: Dict[str, List[torch.Tensor]] = {} + self.B_buffer: Dict[str, List[torch.Tensor]] = {} + + # Lora uid -> buffer idx in memory pool + self.uid_to_buffer_id: Dict[Optional[str], int] = {} + + # Buffer idx -> lora uid in memory pool + # All uids are initialized as `EmptySlot` for empty buffer slots + # Here we don't initialize to None since None is a valid uid + self.buffer_id_to_uid: List[Union[str, None, EmptySlot]] = [ + EMPTY_SLOT + ] * self.max_loras_per_batch + + self.init_buffers(base_model) + + def can_support(self, config: Union[LoRAConfig, Iterable[LoRAConfig]]) -> bool: + """ + Check if the memory pool can support the given LoRA adapters. + """ + + def _can_support(config: LoRAConfig) -> bool: + """ + Check if the memory pool can support a single LoRA adapter. + """ + if config.r > self.max_lora_rank: + return False + target_module_names = get_normalized_target_modules(config.target_modules) + return target_module_names.issubset(self.target_modules) + + if isinstance(config, LoRAConfig): + return _can_support(config) + else: + return all(_can_support(x) for x in config) + + def get_lora_A_shape( + self, + module_name: str, + base_model: torch.nn.Module, + max_lora_dim: int, + layer_idx: int, + ) -> Tuple[int]: + """ + Given a module_name (might be a stacked name), return the hidden dims of modules' input and output. + """ + input_dim, _ = get_hidden_dim( + module_name, self.base_hf_config, base_model, layer_idx + ) + c = get_stacked_multiply(module_name) + if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES: + input_dim = divide(input_dim, self.tp_size) + return ( + self.max_loras_per_batch, + max_lora_dim * c, + input_dim, + ) + + def get_lora_B_shape( + self, + module_name: str, + base_model: torch.nn.Module, + max_lora_dim: int, + layer_idx: int, + ) -> Tuple[int]: + """ + Given a module_name (might be a stacked name), return the hidden dims of modules' input and output. + """ + _, output_dim = get_hidden_dim( + module_name, self.base_hf_config, base_model, layer_idx + ) + if self.tp_size > 1 and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES: + output_dim = divide(output_dim, self.tp_size) + return ( + self.max_loras_per_batch, + output_dim, + max_lora_dim, + ) + + def init_buffers(self, base_model: torch.nn.Module): + device = next(base_model.parameters()).device + + def init_buffer( + buffer: Dict[str, List[torch.Tensor]], + target_modules: Set[str], + get_lora_shape_fn: Callable[[str, torch.nn.Module, int, int], Tuple[int]], + ): + for module_name in target_modules: + buffer[module_name] = [ + torch.empty( + get_lora_shape_fn( + module_name, + base_model, + self.max_lora_rank, + idx, + ), + dtype=self.dtype, + device=device, + ) + for idx in range(self.num_layer) + ] + + init_buffer( + self.A_buffer, + self.target_modules, + self.get_lora_A_shape, + ) + + init_buffer( + self.B_buffer, + self.target_modules, + self.get_lora_B_shape, + ) + + def prepare_lora_batch( + self, + cur_uids: Set[Optional[str]], + lora_adapters: Dict[str, LoRAAdapter], + lora_modules: List[Dict[str, BaseLayerWithLoRA]], + lora_refs: Dict[str, LoRARef], + ): + def get_available_buffer_slot(): + for buffer_id in range(self.max_loras_per_batch): + # Prioritize empty slots + if self.buffer_id_to_uid[buffer_id] == EMPTY_SLOT: + return buffer_id + + for buffer_id in range(self.max_loras_per_batch): + uid = self.buffer_id_to_uid[buffer_id] + + # Evict unneeded lora + if uid not in cur_uids: + # Skip pinned LoRAs + # TODO (lifuhuang): we might consider supporting pinning base model (uid == None) in the future. + if uid is not None: + lora_ref = lora_refs.get(uid) + if lora_ref is not None and lora_ref.pinned: + continue + + self.uid_to_buffer_id.pop(uid) + logger.debug(f"Evicting LoRA {uid} from buffer slot {buffer_id}.") + self.buffer_id_to_uid[buffer_id] = EMPTY_SLOT + return buffer_id + + raise ValueError( + "No available buffer slots found. Please ensure the number of active loras is less than max_loras_per_batch." + ) + + for uid in cur_uids: + if uid not in self.uid_to_buffer_id: + buffer_id = get_available_buffer_slot() + lora_adapter = lora_adapters.get(uid, None) + self.load_lora_weight_to_buffer( + uid, buffer_id, lora_adapter, lora_modules + ) + self.uid_to_buffer_id[uid] = buffer_id + self.buffer_id_to_uid[buffer_id] = uid + + def load_lora_weight_to_buffer( + self, + uid: str, + buffer_id: int, + lora_adapter: LoRAAdapter, + lora_modules: List[Dict[str, BaseLayerWithLoRA]], + ): + def load_lora_weight_tensor( + buffer_view: torch.Tensor, weight: Optional[torch.Tensor] + ): + if weight is None: + # If the particular weight is not present in the adapter, we initialize the buffer to zero + # to avoid contamination from the residual weight of the evicted adapters. + buffer_view.zero_() + else: + assert ( + buffer_view.shape == weight.shape + ), f"LoRA buffer shape {buffer_view.shape} does not match weight shape {weight.shape}." + buffer_view.copy_(weight) + + if uid is None: + for i in range(self.num_layer): + for k in self.A_buffer.keys(): + self.A_buffer[k][i][buffer_id] = 0 + return + + assert lora_adapter is not None + lora_rank = lora_adapter.config.r + for layer_id in range(self.num_layer): + layer_weights = lora_adapter.layers[layer_id].weights + temp_A_buffer: Dict[str, Optional[torch.Tensor]] = { + target_module: None for target_module in self.A_buffer + } + temp_B_buffer: Dict[str, Optional[torch.Tensor]] = { + target_module: None for target_module in self.B_buffer + } + for name, weights in layer_weights.items(): + target_module = get_target_module_name(name, self.target_modules) + if "lora_A" in name: + temp_A_buffer[target_module] = weights + else: + temp_B_buffer[target_module] = weights + + if self.tp_size > 1: + cur_layer_modules = lora_modules[layer_id] + for module_name, module in cur_layer_modules.items(): + target_module = get_target_module_name( + module_name, self.target_modules + ) + + if temp_A_buffer[target_module] is None: + # Skip weight slicing if the weight is not present in the adapter + continue + + temp_A_buffer[target_module] = module.slice_lora_a_weights( + temp_A_buffer[target_module], self.tp_rank + ) + temp_B_buffer[target_module] = module.slice_lora_b_weights( + temp_B_buffer[target_module], self.tp_rank + ) + + for name, weights in temp_A_buffer.items(): + c = get_stacked_multiply(name) + target_buffer = self.A_buffer[name][layer_id] + buffer_view = target_buffer[buffer_id, : lora_rank * c, :] + load_lora_weight_tensor(buffer_view, weights) + + for name, weights in temp_B_buffer.items(): + target_buffer = self.B_buffer[name][layer_id] + buffer_view = target_buffer[buffer_id, :, :lora_rank] + load_lora_weight_tensor(buffer_view, weights) + + def get_tensor( + self, target_module: str, layer_id: int, lora_type: LoRAType + ) -> torch.Tensor: + if lora_type == LoRAType.LORA_A: + return self.A_buffer[target_module][layer_id] + + return self.B_buffer[target_module][layer_id] + + def get_buffer_id(self, lora_uid: str): + return self.uid_to_buffer_id[lora_uid] diff --git a/python/sglang/srt/lora/triton_ops/__init__.py b/python/sglang/srt/lora/triton_ops/__init__.py new file mode 100644 index 000000000..da55e8fd5 --- /dev/null +++ b/python/sglang/srt/lora/triton_ops/__init__.py @@ -0,0 +1,11 @@ +from .gate_up_lora_b import gate_up_lora_b_fwd +from .qkv_lora_b import qkv_lora_b_fwd +from .sgemm_lora_a import sgemm_lora_a_fwd +from .sgemm_lora_b import sgemm_lora_b_fwd + +__all__ = [ + "gate_up_lora_b_fwd", + "qkv_lora_b_fwd", + "sgemm_lora_a_fwd", + "sgemm_lora_b_fwd", +] diff --git a/python/sglang/srt/lora/triton_ops/gate_up_lora_b.py b/python/sglang/srt/lora/triton_ops/gate_up_lora_b.py new file mode 100644 index 000000000..fc4574dd3 --- /dev/null +++ b/python/sglang/srt/lora/triton_ops/gate_up_lora_b.py @@ -0,0 +1,187 @@ +import torch +import triton +import triton.language as tl + +from sglang.srt.lora.utils import LoRABatchInfo + + +@triton.jit +def _gate_up_lora_b_kernel( + # Pointers to matrices + x, + weights, + output, + # Parameters of size + K, # K = R + output_dim, + # Strides + x_stride_0, + x_stride_1, + w_stride_0, + w_stride_1, + w_stride_2, + output_stride_0, + output_stride_1, + # Information on sequence lengths,ranks and weight id + seg_lens, + seg_indptr, + weight_indices, + lora_ranks, + # Meta parameters + BLOCK_S: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + # For fused output scaling + scalings, +): + """ + This kernel packs 2 sgemms (gate/up) into a single kernel. The multiplication + results are accumulated into the output tensor. + + When a sequence's rank is 0, the kernel is essentially a no-op, following + the convention in pytorch where the product of two matrices of shape (m, 0) + and (0, n) is an all-zero matrix of shape (m, n). + + Args: + x (Tensor): The input tensor, which is the result of the LoRA A projection. + Shape: (s, 2 * K), where s is the sum of all sequence lengths in the + batch and K is the maximum LoRA rank. + weights (Tensor): The LoRA B weights for all adapters. + Shape: (num_lora, 2 * output_dim, K). + output (Tensor): The output tensor where the result is stored. + Shape: (s, 2 * output_dim). + """ + # output_dim >> K + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len. + # gate_up_id decides which of gate or up (0: gate, 1: up) + batch_id = tl.program_id(axis=2) + w_index = tl.load(weight_indices + batch_id) + rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel is a no-op. + if rank == 0: + return + + gate_up_id = tl.program_id(axis=1) + pid = tl.program_id(axis=0) + seg_len = tl.load(seg_lens + batch_id) + seg_start = tl.load(seg_indptr + batch_id) + n_start = gate_up_id * output_dim # offset on output dim + scaling = tl.load(scalings + w_index) + + # Adjust K (rank) according to the specific LoRA adapter + K = tl.minimum(K, rank) + + # The tile in output matrix will have (pid_s, pid_n) as id + num_pid_n = tl.cdiv(output_dim, BLOCK_N) + pid_s = pid // num_pid_n + pid_n = pid % num_pid_n + + # Create pointers for the first block of x and weights + # The pointers will be advanced as we move in the K direction + # and accumulate + s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + k_offset = tl.arange(0, BLOCK_K) + + x_ptrs = (x + seg_start * x_stride_0 + (gate_up_id * K) * x_stride_1) + ( + s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + ) + w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < K - k * BLOCK_K) + & (n_offset[None, :] < output_dim), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + + # Store result to output matrix + partial_sum *= scaling + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = (output + seg_start * output_stride_0 + n_start * output_stride_1) + ( + s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + ) + output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < output_dim) + partial_sum += tl.load(output_ptr, mask=output_mask) + tl.store(output_ptr, partial_sum, mask=output_mask) + + +def gate_up_lora_b_fwd( + x: torch.Tensor, + gate_up_lora_b: torch.Tensor, + batch_info: LoRABatchInfo, + output_dim: int, + base_output: torch.Tensor = None, +) -> torch.Tensor: + + # x: (s, 2 * r) + # gate_up_lora_b: (num_lora, 2 * output_dim, r) + # output: (s, 2 * output_dim) + + # Compute lora_output with shape (s, output_dim) as follows: + # lora_output[:, :output_dim] = sgemm(x[:, :r], gate_up_lora_b[:, :output_dim, :]) + # lora_output[:, output_dim:] + # = sgemm(x[:, r:], gate_up_lora_b[:, output_dim:, :]) + + # Get dims + s = x.shape[0] + input_dim = x.shape[1] + r = gate_up_lora_b.shape[-1] + assert input_dim == 2 * r + + BLOCK_S = 16 + BLOCK_R = 16 + BLOCK_OUT = 64 + + grid_b = ( + triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(output_dim, BLOCK_OUT), + 2, # this dimension decides current block computes on gate or up proj + batch_info.bs, + ) + + if base_output is None: + output = torch.zeros((s, 2 * output_dim), device=x.device, dtype=x.dtype) + else: + output = base_output + + _gate_up_lora_b_kernel[grid_b]( + x, + gate_up_lora_b, + output, + r, + output_dim, + x.stride(0), + x.stride(1), + gate_up_lora_b.stride(0), + gate_up_lora_b.stride(1), + gate_up_lora_b.stride(2), + output.stride(0), + output.stride(1), + batch_info.seg_lens, + batch_info.seg_indptr, + batch_info.weight_indices, + batch_info.lora_ranks, + BLOCK_S, + BLOCK_OUT, + BLOCK_R, + batch_info.scalings, + ) + + return output diff --git a/python/sglang/srt/lora/triton_ops/qkv_lora_b.py b/python/sglang/srt/lora/triton_ops/qkv_lora_b.py new file mode 100644 index 000000000..1d6663dbe --- /dev/null +++ b/python/sglang/srt/lora/triton_ops/qkv_lora_b.py @@ -0,0 +1,198 @@ +import torch +import triton +import triton.language as tl + +from sglang.srt.lora.utils import LoRABatchInfo + + +@triton.jit +def _qkv_lora_b_kernel( + # Pointers to matrices + x, + weights, + output, + # Parameters of size + K, # K = R + max_qkv_out_dim, # max(output_q_dim, output_kv_dim) + # Strides + x_stride_0, + x_stride_1, + w_stride_0, + w_stride_1, + w_stride_2, + output_stride_0, + output_stride_1, + # Information on sequence lengths and weight id + seg_lens, + seg_indptr, + weight_indices, + lora_ranks, + # Offsets of q/k/v slice on output dimension + n_offs, + # Meta parameters + BLOCK_S: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + # For fused output scaling + scalings, +): + """ + This kernel packs 3 sgemms (q/k/v) into a single kernel. The multiplication + results are accumulated into the output tensor. + + When a sequence's rank is 0, the kernel is essentially a no-op, following + the convention in pytorch where the product of two matrices of shape (m, 0) + and (0, n) is an all-zero matrix of shape (m, n). + + Args: + x (Tensor): The input tensor, which is the result of the LoRA A projection. + Shape: (s, 3 * K), where s is the sum of all sequence lengths in the + batch and K is the maximum LoRA rank. The second dimension is partitioned + for Q, K, and V. + weights (Tensor): The LoRA B weights for all adapters. + Shape: (num_lora, N_Q + 2 * N_KV, K). + output (Tensor): The output tensor where the result is stored. + Shape: (s, N_Q + 2 * N_KV). + """ + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len. + # qkv_id decides which of q,k,v to compute (0: q, 1: k, 2: v) + batch_id = tl.program_id(axis=2) + w_index = tl.load(weight_indices + batch_id) + rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel is a no-op. + if rank == 0: + return + + qkv_id = tl.program_id(axis=1) + pid = tl.program_id(axis=0) + seg_len = tl.load(seg_lens + batch_id) + seg_start = tl.load(seg_indptr + batch_id) + n_start = tl.load(n_offs + qkv_id) + n_size = tl.load(n_offs + qkv_id + 1) - n_start + scaling = tl.load(scalings + w_index) + # Adjust K (rank) according to the specific LoRA adapter + K = tl.minimum(K, rank) + + # The tile in output matrix will have (pid_s, pid_n) as id + num_pid_n = tl.cdiv(max_qkv_out_dim, BLOCK_N) + pid_s = pid // num_pid_n + pid_n = pid % num_pid_n + + # Create pointers for the first block of x and weights[batch_id][n_start: n_end][:] + # The pointers will be advanced as we move in the K direction + # and accumulate + s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + k_offset = tl.arange(0, BLOCK_K) + + x_ptrs = (x + seg_start * x_stride_0 + (qkv_id * K) * x_stride_1) + ( + s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + ) + w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < K - k * BLOCK_K) & (n_offset[None, :] < n_size), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + + # Store result to output matrix + partial_sum *= scaling + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = (output + seg_start * output_stride_0 + n_start * output_stride_1) + ( + s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + ) + output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < n_size) + partial_sum += tl.load(output_ptr, mask=output_mask) + tl.store(output_ptr, partial_sum, mask=output_mask) + + +def qkv_lora_b_fwd( + x: torch.Tensor, + qkv_lora_b: torch.Tensor, + batch_info: LoRABatchInfo, + output_offset: torch.Tensor, + max_qkv_out_dim: int, + base_output: torch.Tensor = None, +) -> torch.Tensor: + + # x: (s, 3 * r) + # qkv_lora_b: (num_lora, output_dim_q + 2 * output_dim_kv, r) + # output_offset = [0, output_dim_q, output_dim_q + output_dim_kv, + # output_dim_q + 2 * output_dim_kv] + # max_qkv_out_dim = max(output_dim_q, output_dim_kv) + # output: (s, output_dim_q + 2 * output_dim_kv) + + # Compute lora_output with shape (s, output_dim) as follows: + # lora_output[:, :output_dim_q] = sgemm(x[:, :r], qkv_lora_b[:, :outptu_dim_q, :]) + # lora_output[:, output_dim_q: output_dim_q + output_dim_kv] + # = sgemm(x[:, r: 2 * r], qkv_lora_b[:, outptu_dim_q: output_dim_q + output_dim_kv, :]) + # lora_output[:, output_dim_q + output_dim_kv: ] + # = sgemm(x[:, 2 * r: , qkv_lora_b[:, output_dim_q + output_dim_kv: , :]) + + # Get dims + s = x.shape[0] + input_dim = x.shape[1] + r = qkv_lora_b.shape[-1] + output_dim = qkv_lora_b.shape[-2] + assert input_dim == 3 * r + assert output_offset.shape[0] == 4 + + BLOCK_S = 16 + BLOCK_R = 16 + BLOCK_OUT = 64 + + grid_b = ( + triton.cdiv(batch_info.max_len, BLOCK_S) + * triton.cdiv(max_qkv_out_dim, BLOCK_OUT), + 3, # this dimension decides current block computes on q, k or v + batch_info.bs, + ) + + if base_output is None: + output = torch.zeros((s, output_dim), device=x.device, dtype=x.dtype) + else: + output = base_output + + _qkv_lora_b_kernel[grid_b]( + x, + qkv_lora_b, + output, + r, + max_qkv_out_dim, + x.stride(0), + x.stride(1), + qkv_lora_b.stride(0), + qkv_lora_b.stride(1), + qkv_lora_b.stride(2), + output.stride(0), + output.stride(1), + batch_info.seg_lens, + batch_info.seg_indptr, + batch_info.weight_indices, + batch_info.lora_ranks, + output_offset, + BLOCK_S, + BLOCK_OUT, + BLOCK_R, + batch_info.scalings, + ) + + return output diff --git a/python/sglang/srt/lora/triton_ops/sgemm_lora_a.py b/python/sglang/srt/lora/triton_ops/sgemm_lora_a.py new file mode 100644 index 000000000..dded64bcf --- /dev/null +++ b/python/sglang/srt/lora/triton_ops/sgemm_lora_a.py @@ -0,0 +1,170 @@ +import torch +import triton +import triton.language as tl + +from sglang.srt.lora.utils import LoRABatchInfo + + +@triton.jit +def _sgemm_lora_a_kernel( + # Pointers to matrices + x, + weights, + output, + # Matrix dimensions + N, # stack_num * r + K, # input_dim + stack_num, + # Strides + x_stride_0, + x_stride_1, + w_stride_0, + w_stride_1, + w_stride_2, + output_stride_0, + output_stride_1, + # Information on sequence lengths,ranks and weight id + seg_lens, + seg_indptr, + weight_indices, + lora_ranks, + # Meta parameters + BLOCK_S: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """ + Computes a segmented batched matrix multiplication for the LoRA A matrix. + + The kernel ensures that output[seg_start:seg_start + seg_len, :rank * stack_num] + stores the product of the input `x` and the LoRA weights for the corresponding + sequence. This implies that when rank is 0, the kernel is essentially a no-op, + as output[seg_start:seg_start + seg_len, :0] is trivially correct (empty). + + Args: + x (torch.Tensor): The input activations tensor of shape `(s, K)`, where `s` + is the sum of all sequence lengths in the batch. + weights (torch.Tensor): The LoRA 'A' weights for all available adapters, + with shape `(num_lora, N, K)`. + output (torch.Tensor): The output tensor of shape `(s, N)`. + """ + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len + batch_id = tl.program_id(axis=1) + w_index = tl.load(weight_indices + batch_id) + rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel becomes a no-op as the output is always trivially correct. + if rank == 0: + return + + pid = tl.program_id(axis=0) + seg_start = tl.load(seg_indptr + batch_id) + seg_len = tl.load(seg_lens + batch_id) + + # Adjust N (stack_num * max_rank) according to the specific LoRA adapter + N = tl.minimum(N, rank * stack_num) + + # The tile in output matrix will have (pid_s, pid_n) as id + num_pid_n = tl.cdiv(N, BLOCK_N) + pid_s = pid // num_pid_n + pid_n = pid % num_pid_n + + # Create pointers for the first block of x and weights[batch_id] + # The pointers will be advanced as we move in the K direction + # and accumulate + s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + k_offset = tl.arange(0, BLOCK_K) + x_ptrs = (x + seg_start * x_stride_0) + ( + s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + ) + w_ptrs = (weights + w_index * w_stride_0) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < K - k * BLOCK_K) & (n_offset[None, :] < N), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + + # Store result to output matrix + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = (output + seg_start * output_stride_0) + ( + s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + ) + output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < N) + tl.store(output_ptr, partial_sum, mask=output_mask) + + +def sgemm_lora_a_fwd( + x: torch.Tensor, + weights: torch.Tensor, + batch_info: LoRABatchInfo, + stack_num: int = 1, +) -> torch.Tensor: + # x: (s, input_dim) + # weights: (num_lora, stack_num * r, input_dim) + # output: (s, stack_num * r) + # stack_num: run_qkv_lora: 3, run_gate_up_lora: 2 + # when called by run_qkv_lora, the weights.shape[-2] will be 3 * r + # input_dim is much larger than r + + assert x.is_contiguous() + assert weights.is_contiguous() + assert len(x.shape) == 2 + assert len(weights.shape) == 3 + + S = x.shape[0] + R = weights.shape[-2] + K = weights.shape[-1] + assert x.shape[-1] == K + + # Block shapes + BLOCK_S = 16 + BLOCK_K = 256 + BLOCK_R = 16 + + grid = ( + triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(R, BLOCK_R), + batch_info.bs, + ) + + output = torch.empty((S, R), device=x.device, dtype=x.dtype) + _sgemm_lora_a_kernel[grid]( + x, + weights, + output, + R, + K, + stack_num, + x.stride(0), + x.stride(1), + weights.stride(0), + weights.stride(1), + weights.stride(2), + output.stride(0), + output.stride(1), + batch_info.seg_lens, + batch_info.seg_indptr, + batch_info.weight_indices, + batch_info.lora_ranks, + BLOCK_S, + BLOCK_R, + BLOCK_K, + ) + return output diff --git a/python/sglang/srt/lora/triton_ops/sgemm_lora_b.py b/python/sglang/srt/lora/triton_ops/sgemm_lora_b.py new file mode 100644 index 000000000..357d32805 --- /dev/null +++ b/python/sglang/srt/lora/triton_ops/sgemm_lora_b.py @@ -0,0 +1,176 @@ +import torch +import triton +import triton.language as tl + +from sglang.srt.lora.utils import LoRABatchInfo + + +@triton.jit +def _sgemm_lora_b_kernel( + # Pointers to matrices + x, + weights, + output, + # Matrix dimensions + N, # output_dim + K, # r + # Strides + x_stride_0, + x_stride_1, + w_stride_0, + w_stride_1, + w_stride_2, + output_stride_0, + output_stride_1, + # Information on sequence lengths and weight id + seg_lens, + seg_indptr, + weight_indices, + lora_ranks, + # Meta parameters + BLOCK_S: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + # For fused output scaling + scalings, +): + """ + Computes a segmented batched matrix multiplication for the LoRA B matrix + and adds the result to the output in-place. + + When a sequence's rank is 0, the kernel is essentially a no-op, following + the convention in pytorch where the product of two matrices of shape (m, 0) + and (0, n) is an all-zero matrix of shape (m, n). + + Args: + x (torch.Tensor): The intermediate tensor from the LoRA 'A' multiplication, + of shape `(s, K)`, where `s` is the total number of tokens. + weights (torch.Tensor): The LoRA 'B' weights for all available adapters, + with shape `(num_lora, N, K)`. + output (torch.Tensor): The output tensor of shape `(s, N)`. This can be + the base model's output for a fused add operation. + """ + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len + batch_id = tl.program_id(axis=1) + w_index = tl.load(weight_indices + batch_id) + rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel is a no-op. + if rank == 0: + return + + pid = tl.program_id(axis=0) + seg_len = tl.load(seg_lens + batch_id) + seg_start = tl.load(seg_indptr + batch_id) + scaling = tl.load(scalings + w_index) + # Adjust K (rank) according to the specific LoRA adapter + K = tl.minimum(K, rank) + + # The tile in output matrix will have (pid_s, pid_n) as id + num_pid_n = tl.cdiv(N, BLOCK_N) + pid_s = pid // num_pid_n + pid_n = pid % num_pid_n + + # Create pointers for the first block of x and weights[batch_id] + # The pointers will be advanced as we move in the K direction + # and accumulate + s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + k_offset = tl.arange(0, BLOCK_K) + x_ptrs = (x + seg_start * x_stride_0) + ( + s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + ) + w_ptrs = (weights + w_index * w_stride_0) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < K - k * BLOCK_K), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + + # Store result to output matrix + partial_sum *= scaling + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = (output + seg_start * output_stride_0) + ( + s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + ) + output_mask = s_offset[:, None] < seg_len + partial_sum += tl.load(output_ptr, mask=output_mask) + tl.store(output_ptr, partial_sum, mask=output_mask) + + +def sgemm_lora_b_fwd( + x: torch.Tensor, + weights: torch.Tensor, + batch_info: LoRABatchInfo, + base_output: torch.Tensor = None, +) -> torch.Tensor: + # x: (s, max_r) + # weights: (num_lora, output_dim, max_r) + # output: (s, output_dim) + # output_dim is much larger than max_r + + assert x.is_contiguous() + assert weights.is_contiguous() + assert len(x.shape) == 2 + assert len(weights.shape) == 3 + + S = x.shape[0] + N = weights.shape[-2] + R = weights.shape[-1] + assert x.shape[-1] == R + + # Block shapes + BLOCK_S = 16 + BLOCK_R = 16 + BLOCK_N = 256 + + grid = ( + triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(N, BLOCK_N), + batch_info.bs, + ) + + if base_output is None: + output = torch.zeros((S, N), device=x.device, dtype=x.dtype) + else: + output = base_output + + _sgemm_lora_b_kernel[grid]( + x, + weights, + output, + N, + R, + x.stride(0), + x.stride(1), + weights.stride(0), + weights.stride(1), + weights.stride(2), + output.stride(0), + output.stride(1), + batch_info.seg_lens, + batch_info.seg_indptr, + batch_info.weight_indices, + batch_info.lora_ranks, + BLOCK_S, + BLOCK_N, + BLOCK_R, + batch_info.scalings, + ) + return output diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py new file mode 100644 index 000000000..459c943b7 --- /dev/null +++ b/python/sglang/srt/lora/utils.py @@ -0,0 +1,143 @@ +import re +from dataclasses import dataclass +from enum import Enum +from typing import Iterable, Optional, Set, Tuple + +import torch + +from sglang.srt.hf_transformers_utils import AutoConfig + + +@dataclass +class LoRABatchInfo: + # The forward mode is using CUDA Graph. + use_cuda_graph: bool + + # Batch size + bs: int + + # Number of segments. For triton backend, it is equal to batch size. + num_segments: int + + # Indice pointers of each segment in shape (num_segments + 1, ) + seg_indptr: torch.Tensor + + # The index of lora adapter used by each segment, in shape (num_segments,) + weight_indices: torch.Tensor + + # ranks of each lora adapter, in shape (lora_num,) + lora_ranks: torch.Tensor + + # scaling of each lora adapter, in shape (lora_num,) + scalings: torch.Tensor + + # Lengths of each segments in shape (num_segments,) + seg_lens: Optional[torch.Tensor] + + # Maximum segment length of current batch + max_len: Optional[int] + + # The logical (re)ordering of input rows (tokens), in shape (num_tokens,) + permutation: Optional[torch.Tensor] + + +class LoRAType(Enum): + LORA_A = 0 + LORA_B = 1 + + +def get_layer_id(name: str) -> int: + """ + Extract integer id of layer from its name in string. + """ + match = re.search(r"layers\.(\d+)\.", name) + if match is None: + return None + return int(match.group(1)) + + +def get_hidden_dim( + module_name: str, config: AutoConfig, base_model: torch.nn.Module, layer_idx: int +) -> Tuple[int]: + """ + Given a module_name (might be a stacked name), return the hidden dims of modules' input and output. + """ + + if hasattr(base_model, "get_hidden_dim"): + return base_model.get_hidden_dim(module_name, layer_idx) + else: + """ + WARNING: get_hidden_dim() is not defined, + which is used to get the hidden dim for different lora modules + Use the default one, but please check if it is correct for your model. + Please implement the function in the model class if it is not. + You can reference this function in llama.py. + """ + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + if module_name == "qkv_proj": + return config.hidden_size, head_dim * ( + config.num_attention_heads + config.num_key_value_heads * 2 + ) + elif module_name == "o_proj": + return ( + head_dim * config.num_attention_heads, + config.hidden_size, + ) + elif module_name == "gate_up_proj": + return config.hidden_size, config.intermediate_size * 2 + elif module_name == "down_proj": + return config.intermediate_size, config.hidden_size + else: + raise NotImplementedError() + + +def get_normalized_target_modules( + target_modules: Iterable[str], +) -> set[str]: + """ + Mapping a list of target module name to names of the normalized LoRA weights. + """ + params_mapping = { + "q_proj": "qkv_proj", + "k_proj": "qkv_proj", + "v_proj": "qkv_proj", + "gate_proj": "gate_up_proj", + "up_proj": "gate_up_proj", + } + + result = set() + for name in target_modules: + normalized_name = params_mapping.get(name, name) + result.add(normalized_name) + return result + + +def get_stacked_multiply(module_name: str) -> int: + """ + Mapping a lora module name to its magnification at output dimension + """ + stacked_rank = { + "qkv_proj": 3, + "gate_up_proj": 2, + } + return stacked_rank[module_name] if module_name in stacked_rank else 1 + + +def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> str: + """ + Get the target module name in target_modules that can match full_module_name. + + If there is a target module name in target_modules that can match full_module_name, return this name + Else raise ValueError. + """ + for target_module in target_modules: + if target_module in full_module_name: + return target_module + raise ValueError( + f"Cannot find target module name for {full_module_name} in {target_modules}" + ) + + +ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj"] diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py new file mode 100644 index 000000000..f9d45b2f7 --- /dev/null +++ b/python/sglang/srt/managers/cache_controller.py @@ -0,0 +1,886 @@ +from __future__ import annotations + +""" +Copyright 2023-2025 SGLang 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 logging +import math +import threading +import time +from queue import Empty, Full, PriorityQueue, Queue +from typing import TYPE_CHECKING, List, NamedTuple, Optional, Set, Tuple + +import torch + +from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig + +if TYPE_CHECKING: + from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator + from sglang.srt.mem_cache.memory_pool_host import HostKVCache + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from sglang.srt.layers.dp_attention import ( + get_attention_dp_rank, + get_attention_tp_rank, + get_attention_tp_size, + is_dp_attention_enabled, +) +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool + +logger = logging.getLogger(__name__) + + +class LayerLoadingEvent: + def __init__(self, num_layers: int): + self._num_layers = num_layers + self.load_events = [torch.cuda.Event() for _ in range(num_layers)] + self.start_event = torch.cuda.Event() # start event on controller stream + + def complete(self, layer_index: int): + assert 0 <= layer_index < self._num_layers + self.load_events[layer_index].record() + + def wait(self, layer_index: int): + torch.cuda.current_stream().wait_event(self.load_events[layer_index]) + + @property + def finish_event(self): + return self.load_events[-1] + + +class LayerDoneCounter: + def __init__(self, num_layers: int): + self.num_layers = num_layers + # extra producer and consumer counters for overlap mode + self.num_counters = 3 + self.events = [LayerLoadingEvent(num_layers) for _ in range(self.num_counters)] + self.producer_index = -1 + self.consumer_index = -1 + + def update_producer(self): + self.producer_index = (self.producer_index + 1) % self.num_counters + assert self.events[ + self.producer_index + ].finish_event.query(), ( + "Producer finish event should be ready before being reused." + ) + return self.producer_index + + def set_consumer(self, index: int): + self.consumer_index = index + + def wait_until(self, threshold: int): + if self.consumer_index < 0: + return + self.events[self.consumer_index].wait(threshold) + + def reset(self): + self.producer_index = -1 + self.consumer_index = -1 + + +class CacheOperation: + + counter = 0 + + def __init__( + self, + host_indices: torch.Tensor, + device_indices: torch.Tensor, + node_id: int, + priority: Optional[int] = None, + ): + self.host_indices = host_indices + self.device_indices = device_indices + self.node_ids = [node_id] + self.data = None + + self.id = CacheOperation.counter + CacheOperation.counter += 1 + # default priority is the order of creation + self.priority = priority if priority is not None else self.id + + @staticmethod + def merge_ops(ops: List[CacheOperation]) -> CacheOperation: + assert len(ops) > 0 + if len(ops) == 1: + return ops[0] + + host_indices = torch.cat([op.host_indices for op in ops]) + device_indices = torch.cat([op.device_indices for op in ops]) + node_ids = [] + priority = min(op.priority for op in ops) + for op in ops: + node_ids.extend(op.node_ids) + merged_op = CacheOperation(host_indices, device_indices, -1, priority) + merged_op.node_ids = node_ids + return merged_op + + def __lt__(self, other: CacheOperation): + return self.priority < other.priority + + +class HiCacheAck(NamedTuple): + start_event: torch.cuda.Event + finish_event: torch.cuda.Event + node_ids: List[int] + + +class TransferBuffer: + """ + Overlapping buffer preparation and transfer operations to improve throughput. + """ + + def __init__( + self, stop_event, buffer_count: int = 3, max_buffer_size: int = 1024 + ) -> None: + self.stop_event = stop_event + self.buffers = Queue(maxsize=buffer_count) + # todo: adjust the buffer size based on throughput profile of the system + self.max_buffer_size = max_buffer_size + + def full(self) -> bool: + return self.buffers.full() + + def empty(self) -> bool: + return self.buffers.empty() + + def put(self, item, block=True, timeout=1) -> None: + while not self.stop_event.is_set(): + try: + self.buffers.put(item, block=block, timeout=timeout) + break + except Full: + if not block: + break + continue + except Exception as e: + logger.error(e) + + def get(self, block=True, timeout=1) -> Optional[CacheOperation]: + try: + return self.buffers.get(block=block, timeout=timeout) + except Empty: + return None + except Exception as e: + logger.error(e) + + def clear(self): + self.buffers.queue.clear() + + +class StorageOperation: + counter = 0 + + def __init__( + self, + host_indices: torch.Tensor, + token_ids: List[int], + last_hash: Optional[str] = None, + hash_value: Optional[List[str]] = None, + ): + self.host_indices = host_indices + self.token_ids = token_ids + self.last_hash = last_hash + self.completed_tokens = 0 + self.hash_value = hash_value if hash_value is not None else [] + + self.id = StorageOperation.counter + StorageOperation.counter += 1 + + def __lt__(self, other: "StorageOperation"): + return self.id < other.id + + +class PrefetchOperation(StorageOperation): + def __init__( + self, + request_id: str, + host_indices: torch.Tensor, + token_ids: List[int], + last_hash: Optional[str] = None, + ): + self.request_id = request_id + + self._lock = threading.Lock() + self._terminated_flag = False + self.start_time = time.monotonic() + + super().__init__(host_indices, token_ids, last_hash) + + def increment(self, num_tokens: int): + with self._lock: + if self._terminated_flag: + return False + self.completed_tokens += num_tokens + return True + + def mark_terminate(self): + with self._lock: + self._terminated_flag = True + + def is_terminated(self) -> bool: + return self._terminated_flag + + +class HiCacheController: + + def __init__( + self, + token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, + mem_pool_host: HostKVCache, + page_size: int, + tp_group: torch.distributed.ProcessGroup, + load_cache_event: threading.Event, + write_policy: str = "write_through_selective", + io_backend: str = "", + storage_backend: Optional[str] = None, + prefetch_threshold: int = 256, + model_name: Optional[str] = None, + storage_backend_extra_config: Optional[str] = None, + ): + self.mem_pool_device_allocator = token_to_kv_pool_allocator + self.mem_pool_device = token_to_kv_pool_allocator.get_kvcache() + self.mem_pool_host = mem_pool_host + self.write_policy = write_policy + self.page_size = page_size + self.io_backend = io_backend + self.enable_storage = False + + if storage_backend is not None: + self.storage_backend_type = storage_backend + from sglang.srt.mem_cache.hicache_storage import get_hash_str + + self.get_hash_str = get_hash_str + self.storage_config = self._generate_storage_config( + model_name, storage_backend_extra_config + ) + # for MLA models, only one rank needs to backup the KV cache + self.backup_skip = ( + self.storage_config.is_mla_model + # todo: load balancing + and self.storage_config.tp_rank != 0 + ) + + if storage_backend == "file": + from sglang.srt.mem_cache.hicache_storage import HiCacheFile + + self.storage_backend = HiCacheFile(self.storage_config) + elif storage_backend == "nixl": + from sglang.srt.mem_cache.storage.nixl.hicache_nixl import HiCacheNixl + + self.storage_backend = HiCacheNixl() + elif storage_backend == "mooncake": + from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( + MooncakeStore, + ) + + self.storage_backend = MooncakeStore(self.storage_config) + self.storage_backend.register_buffer(self.mem_pool_host.kv_buffer) + assert self.mem_pool_host.layout == "page_first" + elif storage_backend == "hf3fs": + from sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs import ( + HiCacheHF3FS, + ) + + if self.mem_pool_host.layout == "page_first": + bytes_per_page = ( + mem_pool_host.get_ksize_per_token() * mem_pool_host.page_size + ) + elif self.mem_pool_host.layout == "layer_first": + bytes_per_page = ( + mem_pool_host.get_size_per_token() * mem_pool_host.page_size + ) + dtype = mem_pool_host.dtype + self.storage_backend = HiCacheHF3FS.from_env_config( + bytes_per_page, dtype, self.storage_config + ) + else: + raise NotImplementedError( + f"Unsupported storage backend: {storage_backend}" + ) + + self.enable_storage = True + # todo: threshold policy for prefetching + self.prefetch_threshold = max(prefetch_threshold, self.page_size) + self.prefetch_capacity_limit = int( + 0.8 * (self.mem_pool_host.size - self.mem_pool_device.size) + ) + # granularity of batch storage IO operations, in number of pages + self.storage_batch_size = 128 + # tracking the number of tokens locked in prefetching, updated by the main scheduler thread + self.prefetch_tokens_occupied = 0 + + # create a new communication group for synchronizing storage operations across TP workers + self.tp_world_size = torch.distributed.get_world_size(group=tp_group) + if self.tp_world_size > 1: + group_ranks = torch.distributed.get_process_group_ranks(tp_group) + self.prefetch_tp_group = torch.distributed.new_group( + group_ranks, backend="gloo" + ) + + # Select the get and set functions + self.page_get_func = self._generic_page_get + self.page_set_func = self._generic_page_set + self.batch_exists_func = self.storage_backend.batch_exists + self.is_3fs_zerocopy = ( + self.storage_backend_type == "hf3fs" + and self.mem_pool_host.layout == "page_first" + ) + if self.storage_backend_type == "mooncake": + self.page_get_func = self._mooncake_page_get + self.page_set_func = self._mooncake_page_set + elif self.is_3fs_zerocopy: + self.page_get_func = self._3fs_zero_copy_page_get + self.page_set_func = self._3fs_zero_copy_page_set + self.batch_exists_func = self._3fs_zero_copy_batch_exists + + self.device = self.mem_pool_device.device + self.layer_num = self.mem_pool_device.layer_num + self.layer_done_counter = LayerDoneCounter(self.layer_num) + self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter) + + if write_policy not in [ + "write_through", + "write_through_selective", + "write_back", + ]: + raise ValueError(f"Invalid write policy: {write_policy}") + + # self.write_queue = PriorityQueue[CacheOperation]() + self.load_queue: List[CacheOperation] = [] + self.write_queue: List[CacheOperation] = [] + self.ack_load_queue: List[HiCacheAck] = [] + self.ack_write_queue: List[HiCacheAck] = [] + + self.stop_event = threading.Event() + self.write_buffer = TransferBuffer(self.stop_event) + self.load_buffer = TransferBuffer( + self.stop_event, buffer_count=10, max_buffer_size=100 + ) + + self.write_stream = torch.cuda.Stream() + self.load_stream = torch.cuda.Stream() + + if self.enable_storage: + self.prefetch_thread = threading.Thread( + target=self.prefetch_thread_func, daemon=True + ) + self.backup_thread = threading.Thread( + target=self.backup_thread_func, daemon=True + ) + self.prefetch_queue = Queue() + self.backup_queue = Queue() + + self.prefetch_revoke_queue = Queue() + self.ack_backup_queue = Queue() + self.host_mem_release_queue = Queue() + + self.prefetch_thread.start() + self.backup_thread.start() + + def _generate_storage_config( + self, + model_name: Optional[str] = None, + storage_backend_extra_config: Optional[str] = None, + ): + + if is_dp_attention_enabled(): + self.tp_rank = get_attention_tp_rank() + self.tp_size = get_attention_tp_size() + self.dp_rank = get_attention_dp_rank() + else: + self.tp_rank = get_tensor_model_parallel_rank() + self.tp_size = get_tensor_model_parallel_world_size() + self.dp_rank = 0 + + # Currently, AscendMLAPagedTokenToKVPool is the subclass of MLATokenToKVPool. + is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool) + + # Parse extra config JSON if provided + extra_config = None + if storage_backend_extra_config: + try: + import json + + extra_config = json.loads(storage_backend_extra_config) + except Exception as e: + logger.error(f"Invalid backend extra config JSON: {e}") + + return HiCacheStorageConfig( + tp_rank=self.tp_rank, + tp_size=self.tp_size, + is_mla_model=is_mla_backend, + is_page_first_layout=self.mem_pool_host.layout == "page_first", + model_name=model_name, + extra_config=extra_config, + ) + + def reset(self): + self.stop_event.set() + + self.write_queue.clear() + self.load_queue.clear() + self.write_buffer.clear() + self.load_buffer.clear() + self.ack_write_queue.clear() + self.ack_load_queue.clear() + if self.enable_storage: + self.prefetch_thread.join() + self.backup_thread.join() + self.prefetch_queue.queue.clear() + self.backup_queue.queue.clear() + self.prefetch_revoke_queue.queue.clear() + self.ack_backup_queue.queue.clear() + + self.stop_event.clear() + + if self.enable_storage: + self.prefetch_thread = threading.Thread( + target=self.prefetch_thread_func, daemon=True + ) + self.backup_thread = threading.Thread( + target=self.backup_thread_func, daemon=True + ) + self.prefetch_thread.start() + self.backup_thread.start() + + def write( + self, + device_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = -1, + ) -> Optional[torch.Tensor]: + """ + Back up KV caches from device memory to host memory. + """ + host_indices = self.mem_pool_host.alloc(len(device_indices)) + if host_indices is None: + return None + self.mem_pool_host.protect_write(host_indices) + self.write_queue.append( + CacheOperation(host_indices, device_indices, node_id, priority) + ) + self.start_writing() + return host_indices + + def start_writing(self) -> None: + if len(self.write_queue) == 0: + return + + op = CacheOperation.merge_ops(self.write_queue) + host_indices, device_indices = self.move_indices(op) + self.write_queue.clear() + + start_event = torch.cuda.Event() + finish_event = torch.cuda.Event() + + start_event.record() + with torch.cuda.stream(self.write_stream): + start_event.wait(self.write_stream) + self.mem_pool_host.backup_from_device_all_layer( + self.mem_pool_device, host_indices, device_indices, self.io_backend + ) + self.mem_pool_host.complete_io(op.host_indices) + finish_event.record() + # NOTE: We must save the host indices and device indices here, + # this is because we need to guarantee that these tensors are + # still alive when the write stream is executing. + if host_indices.is_cuda: + host_indices.record_stream(self.write_stream) + if device_indices.is_cuda: + device_indices.record_stream(self.write_stream) + + self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + + def load( + self, + host_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = -1, + ) -> Optional[torch.Tensor]: + """ + Load KV caches from host memory to device memory. + """ + device_indices = self.mem_pool_device_allocator.alloc(len(host_indices)) + if device_indices is None: + return None + self.mem_pool_host.protect_load(host_indices) + self.load_queue.append( + CacheOperation(host_indices, device_indices, node_id, priority) + ) + return device_indices + + def move_indices(self, op: CacheOperation): + host_indices, device_indices = op.host_indices, op.device_indices + # move indices to GPU if using kernels, to host if using direct indexing + if self.io_backend == "kernel": + if not host_indices.is_cuda: + host_indices = host_indices.to(self.device, non_blocking=True) + return host_indices, device_indices + elif self.io_backend == "direct": + device_indices = device_indices.cpu() + host_indices, idx = host_indices.sort() + return host_indices, device_indices.index_select(0, idx) + else: + raise ValueError(f"Unsupported io backend") + + def start_loading(self) -> int: + if len(self.load_queue) == 0: + return -1 + + producer_id = self.layer_done_counter.update_producer() + op = CacheOperation.merge_ops(self.load_queue) + host_indices, device_indices = self.move_indices(op) + self.load_queue.clear() + producer_event = self.layer_done_counter.events[producer_id] + producer_event.start_event.record() + + with torch.cuda.stream(self.load_stream): + producer_event.start_event.wait(self.load_stream) + for i in range(self.layer_num): + self.mem_pool_host.load_to_device_per_layer( + self.mem_pool_device, + host_indices, + device_indices, + i, + self.io_backend, + ) + producer_event.complete(i) + self.mem_pool_host.complete_io(op.host_indices) + # NOTE: We must save the host indices and device indices here, + # this is because we need to guarantee that these tensors are + # still alive when the load stream is executing. + if host_indices.is_cuda: + host_indices.record_stream(self.load_stream) + if device_indices.is_cuda: + device_indices.record_stream(self.load_stream) + + self.ack_load_queue.append( + HiCacheAck( + start_event=producer_event.start_event, + finish_event=producer_event.finish_event, + node_ids=op.node_ids, + ) + ) + return producer_id + + def evict_device( + self, device_indices: torch.Tensor, host_indices: torch.Tensor + ) -> int: + if self.mem_pool_host.is_synced(host_indices): + self.mem_pool_device_allocator.free(device_indices) + self.mem_pool_host.update_backup(host_indices) + return len(device_indices) + else: + raise ValueError( + f"Inconsistent states: {self.mem_pool_host.get_state(host_indices)}" + ) + + def evict_host(self, host_indices: torch.Tensor, backup_only: bool = True) -> int: + if not backup_only: + raise ValueError("Other eviction policies are not supported yet.") + + if self.mem_pool_host.is_backup(host_indices): + self.mem_pool_host.free(host_indices) + return len(host_indices) + else: + raise ValueError( + f"Inconsistent states: {self.mem_pool_host.get_state(host_indices)}" + ) + + def prefetch( + self, + request_id: str, + host_indices: torch.Tensor, + new_input_tokens: List[int], + last_hash: Optional[str] = None, + ) -> PrefetchOperation: + """ + Prefetch KV caches from storage backend to host memory. + """ + operation = PrefetchOperation( + request_id, host_indices, new_input_tokens, last_hash + ) + self.prefetch_queue.put(operation) + return operation + + def terminate_prefetch(self, operation): + operation.mark_terminate() + return operation.completed_tokens, operation.hash_value + + def append_host_mem_release(self, host_indices: torch.Tensor): + chunks = host_indices.split(self.mem_pool_host.page_size) + for chunk in chunks: + self.host_mem_release_queue.put(chunk) + + def _3fs_zero_copy_batch_exists(self, batch_hashes): + _batch_hashes, _, factor = self.mem_pool_host.get_buffer_with_hash(batch_hashes) + hit_page_num = self.storage_backend.batch_exists(_batch_hashes) // factor + return hit_page_num + + def _3fs_zero_copy_page_get(self, operation, hash_values, host_indices): + hashes, dsts, factor = self.mem_pool_host.get_buffer_with_hash( + hash_values, host_indices + ) + page_data = self.storage_backend.batch_get(hashes, dsts) + if page_data: + inc = self.page_size * len(hashes) // factor + operation.increment(inc) + else: + logger.warning( + f"Prefetch operation {operation.request_id} failed to retrieve page {hashes}." + ) + + def _mooncake_page_get(self, operation, hash_values, host_indices): + key_strs, buffer_ptrs, buffer_sizes = self.mem_pool_host.get_buffer_meta( + hash_values, + host_indices, + self.storage_config.tp_rank, + ) + get_result = self.storage_backend.batch_get( + key_strs, + target_locations=buffer_ptrs, + target_sizes=buffer_sizes, + ) + if get_result != len(hash_values): + logger.warning( + f"Prefetch operation {operation.request_id} failed or partially failed." + ) + if get_result != 0: + operation.increment(get_result * self.page_size) + + def _generic_page_get(self, operation, hash_values, host_indices): + dummy_page_dst = [ + self.mem_pool_host.get_dummy_flat_data_page() for _ in hash_values + ] + page_data = self.storage_backend.batch_get(hash_values, dummy_page_dst) + if page_data is None: + return + for i in range(len(hash_values)): + if page_data[i] is None: + logger.warning( + f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}." + ) + break + # Must set the data before increasing the completed tokens. + # Otherwise this page may be read before being set. + self.mem_pool_host.set_from_flat_data_page( + host_indices[i * self.page_size], + page_data[i], + ) + if not operation.increment(self.page_size): + break # Operation terminated by controller + + def _page_transfer(self, operation): + # Transfer batch by batch + for i in range(0, len(operation.hash_value), self.storage_batch_size): + batch_hashes = operation.hash_value[i : i + self.storage_batch_size] + batch_host_indices = operation.host_indices[ + i * self.page_size : (i + len(batch_hashes)) * self.page_size + ] + prev_completed_tokens = operation.completed_tokens + # Get one batch token, and update the completed_tokens if succeed + self.page_get_func(operation, batch_hashes, batch_host_indices) + # Check termination + if ( + operation.completed_tokens + != prev_completed_tokens + len(batch_hashes) * self.page_size + ): + operation.mark_terminate() + break # Some operations fail or operation terminated by controller + # release pre-allocated memory + self.append_host_mem_release( + operation.host_indices[operation.completed_tokens :] + ) + + def prefetch_io_aux_func(self): + """ + Auxiliary function conducting IO operations for prefetching. + """ + while not self.stop_event.is_set(): + try: + operation = self.prefetch_buffer.get(block=True, timeout=1) + self._page_transfer(operation) + # operation terminated by controller, release pre-allocated memory + self.append_host_mem_release( + operation.host_indices[operation.completed_tokens :] + ) + except Empty: + continue + + def prefetch_rate_limited(self) -> bool: + """ + Rate limit the prefetching operations to avoid overwhelming the storage backend. + """ + # cancel prefetch if too much memory is occupied + if self.prefetch_tokens_occupied >= self.prefetch_capacity_limit: + return True + # todo: more sophisticated rate limiting based on storage backend performance + return False + + def _storage_hit_query(self, operation) -> tuple[list[str], int]: + last_hash = operation.last_hash + tokens_to_fetch = operation.token_ids + + storage_query_count = 0 + hash_value = [] + + for start in range( + 0, len(tokens_to_fetch), self.page_size * self.storage_batch_size + ): + end = min( + start + self.page_size * self.storage_batch_size, len(tokens_to_fetch) + ) + batch_tokens = tokens_to_fetch[start:end] + batch_hashes = [] + for i in range(0, len(batch_tokens), self.page_size): + last_hash = self.get_hash_str( + batch_tokens[i : i + self.page_size], last_hash + ) + batch_hashes.append(last_hash) + hit_page_num = self.batch_exists_func(batch_hashes) + hash_value.extend(batch_hashes[:hit_page_num]) + storage_query_count += hit_page_num * self.page_size + if hit_page_num < len(batch_hashes): + break + return hash_value, storage_query_count + + def prefetch_thread_func(self): + """ + Manage prefetching operations from storage backend to host memory. + """ + self.prefetch_buffer = Queue() + aux_thread = threading.Thread(target=self.prefetch_io_aux_func, daemon=True) + aux_thread.start() + while (not self.stop_event.is_set()) or not self.prefetch_queue.empty(): + try: + operation = self.prefetch_queue.get(block=True, timeout=1) + if operation is None: + continue + + hash_value, storage_hit_count = self._storage_hit_query(operation) + if self.tp_world_size > 1: + storage_hit_count_tensor = torch.tensor( + storage_hit_count, dtype=torch.int + ) + torch.distributed.all_reduce( + storage_hit_count_tensor, + op=torch.distributed.ReduceOp.MIN, + group=self.prefetch_tp_group, + ) + storage_hit_count = storage_hit_count_tensor.item() + + if storage_hit_count < self.prefetch_threshold: + # not to prefetch if not enough benefits + self.prefetch_revoke_queue.put(operation.request_id) + self.append_host_mem_release(operation.host_indices) + logger.debug( + f"Revoking prefetch for request {operation.request_id} due to insufficient hits ({storage_hit_count})." + ) + else: + operation.hash_value = hash_value[ + : (storage_hit_count // self.page_size) + ] + # free the pre-allocated memory for pages that are not hit + self.append_host_mem_release( + operation.host_indices[storage_hit_count:] + ) + operation.host_indices = operation.host_indices[:storage_hit_count] + logger.debug( + f"Prefetching {len(operation.hash_value)} pages for request {operation.request_id}." + ) + self.prefetch_buffer.put(operation) + + except Empty: + continue + + def write_storage( + self, + host_indices: torch.Tensor, + token_ids: List[int], + hash_value: Optional[List[str]] = None, + ) -> int: + """ + Write KV caches from host memory to storage backend. + """ + operation = StorageOperation(host_indices, token_ids, hash_value=hash_value) + self.backup_queue.put(operation) + return operation.id + + # non-zero copy + def _generic_page_set(self, hash_values, host_indices) -> bool: + data = [ + self.mem_pool_host.get_flat_data_page(host_indices[i * self.page_size]) + for i in range(len(hash_values)) + ] + return self.storage_backend.batch_set(hash_values, data) + + # zero copy + def _mooncake_page_set(self, hash_values, host_indices) -> bool: + key_strs, buffer_ptrs, buffer_sizes = self.mem_pool_host.get_buffer_meta( + hash_values, + host_indices, + self.storage_config.tp_rank, + ) + success = self.storage_backend.batch_set( + key_strs, + target_locations=buffer_ptrs, + target_sizes=buffer_sizes, + ) + return success + + # zero copy + def _3fs_zero_copy_page_set(self, hash_values, host_indices) -> bool: + hashes, dsts, _ = self.mem_pool_host.get_buffer_with_hash( + hash_values, host_indices + ) + return self.storage_backend.batch_set(hashes, dsts) + + # Backup batch by batch + def _page_backup(self, operation): + # Backup batch by batch + for i in range(0, len(operation.hash_value), self.storage_batch_size): + batch_hashes = operation.hash_value[i : i + self.storage_batch_size] + batch_host_indices = operation.host_indices[ + i * self.page_size : (i + len(batch_hashes)) * self.page_size + ] + # Set one batch token, and record if success. + # todo: allow partial success + success = self.page_set_func(batch_hashes, batch_host_indices) + if not success: + logger.warning( + f"Write page to storage: {len(batch_hashes)} pages failed." + ) + break + operation.completed_tokens += self.page_size * len(batch_hashes) + + def backup_thread_func(self): + """ + Manage backup operations from host memory to storage backend. + """ + while not self.stop_event.is_set(): + try: + operation = self.backup_queue.get(block=True, timeout=1) + if operation is None: + continue + + if not self.backup_skip: + self._page_backup(operation) + self.ack_backup_queue.put(operation) + + except Empty: + continue diff --git a/python/sglang/srt/managers/configure_logging.py b/python/sglang/srt/managers/configure_logging.py new file mode 100644 index 000000000..0dc78edfa --- /dev/null +++ b/python/sglang/srt/managers/configure_logging.py @@ -0,0 +1,47 @@ +""" +Copyright 2023-2025 SGLang 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. +""" + +""" +Configure the logging settings of a server. + +Usage: +python3 -m sglang.srt.managers.configure_logging --url http://localhost:30000 +""" + +import argparse + +import requests + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--url", type=str, default="http://localhost:30000") + parser.add_argument("--log-requests", action="store_true") + parser.add_argument("--log-requests-level", type=int, default=3) + parser.add_argument( + "--dump-requests-folder", type=str, default="/tmp/sglang_request_dump" + ) + parser.add_argument("--dump-requests-threshold", type=int, default=1000) + args = parser.parse_args() + + response = requests.post( + args.url + "/configure_logging", + json={ + "log_requests": args.log_requests, + "log_requests_level": args.log_requests_level, # Log full requests + "dump_requests_folder": args.dump_requests_folder, + "dump_requests_threshold": args.dump_requests_threshold, + }, + ) + assert response.status_code == 200 diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py new file mode 100644 index 000000000..a7bb6d13a --- /dev/null +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -0,0 +1,390 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""A controller that dispatches requests to multiple data parallel workers.""" + +import faulthandler +import logging +import multiprocessing as mp +import signal +import struct +import sys +import threading +import time +from enum import Enum, auto +from multiprocessing import shared_memory +from typing import Dict, List + +import psutil +import setproctitle +import zmq + +from sglang.srt.layers.dp_attention import compute_dp_attention_world_info +from sglang.srt.managers.io_struct import ( + BlockReqInput, + TokenizedEmbeddingReqInput, + TokenizedGenerateReqInput, +) +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.managers.scheduler import run_scheduler_process +from sglang.srt.managers.utils import DPBalanceMeta +from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter +from sglang.srt.utils import ( + bind_port, + configure_logger, + get_zmq_socket, + kill_itself_when_parent_died, +) +from sglang.utils import get_exception_traceback + +logger = logging.getLogger(__name__) + + +class LoadBalanceMethod(Enum): + """Load balance method.""" + + ROUND_ROBIN = auto() + SHORTEST_QUEUE = auto() + MINIMUM_TOKENS = auto() + + @classmethod + def from_str(cls, method: str): + method = method.upper() + try: + return cls[method] + except KeyError as exc: + raise ValueError(f"Invalid load balance method: {method}") from exc + + +class DataParallelController: + """A controller that dispatches requests to multiple data parallel workers.""" + + def __init__( + self, + server_args: ServerArgs, + port_args: PortArgs, + dp_balance_meta: DPBalanceMeta, + ) -> None: + # for dp balance + self.global_balance_id = 0 + self.balance_meta = dp_balance_meta + + # Parse args + self.max_total_num_tokens = None + self.server_args = server_args + self.port_args = port_args + self.load_balance_method = LoadBalanceMethod.from_str( + server_args.load_balance_method + ) + + # Init inter-process communication + self.context = zmq.Context(1 + server_args.dp_size) + if server_args.node_rank == 0: + self.recv_from_tokenizer = get_zmq_socket( + self.context, zmq.PULL, port_args.scheduler_input_ipc_name, False + ) + + # Dispatch method + self.round_robin_counter = 0 + dispatch_lookup = { + LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler, + LoadBalanceMethod.SHORTEST_QUEUE: self.shortest_queue_scheduler, + LoadBalanceMethod.MINIMUM_TOKENS: self.minimum_tokens_scheduler, + } + self.dispatching = dispatch_lookup[self.load_balance_method] + + # Launch data parallel workers + self.scheduler_procs = [] + self.workers: List[zmq.Socket] = [None] * server_args.dp_size + + if server_args.enable_dp_attention: + dp_port_args = self.launch_dp_attention_schedulers(server_args, port_args) + self.control_message_step = server_args.tp_size + else: + dp_port_args = self.launch_dp_schedulers(server_args, port_args) + self.control_message_step = 1 + + # Only node rank 0 runs the real data parallel controller that dispatches the requests. + if server_args.node_rank == 0: + for dp_rank in range(server_args.dp_size): + self.workers[dp_rank] = get_zmq_socket( + self.context, + zmq.PUSH, + dp_port_args[dp_rank].scheduler_input_ipc_name, + True, + ) + + self.max_req_input_len = None + + def launch_dp_schedulers(self, server_args, port_args): + base_gpu_id = 0 + + threads = [] + sockets = [] + dp_port_args = [] + ready_events = [] + for dp_rank in range(server_args.dp_size): + tmp_port_args = PortArgs.init_new(server_args) + tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name + tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name + dp_port_args.append(tmp_port_args) + + # This port is checked free in PortArgs.init_new. + # We hold it first so that the next dp worker gets a different port + sockets.append(bind_port(tmp_port_args.nccl_port)) + + ready_event = threading.Event() + ready_events.append(ready_event) + + # Create a thread for each worker + thread = threading.Thread( + target=self.launch_tensor_parallel_group_thread, + args=(server_args, tmp_port_args, base_gpu_id, dp_rank, ready_event), + ) + threads.append(thread) + base_gpu_id += server_args.tp_size * server_args.gpu_id_step + + # Free all sockets before starting the threads to launch TP workers + for sock in sockets: + sock.close() + + # Start all threads + for thread in threads: + thread.start() + for event in ready_events: + event.wait() + + return dp_port_args + + def launch_tensor_parallel_group_thread( + self, + server_args: ServerArgs, + port_args: PortArgs, + base_gpu_id: int, + dp_rank: int, + ready_event: threading.Event, + ): + self.launch_tensor_parallel_group(server_args, port_args, base_gpu_id, dp_rank) + ready_event.set() + + # This thread cannot be closed because otherwise the `kill_itself_when_parent_died` + # function in scheduler.py will kill the scheduler. + while True: + time.sleep(30 * 24 * 3600) + + def launch_dp_attention_schedulers(self, server_args, port_args): + self.launch_tensor_parallel_group(server_args, port_args, 0, None) + dp_port_args = [] + for dp_rank in range(server_args.dp_size): + dp_port_args.append(PortArgs.init_new(server_args, dp_rank)) + return dp_port_args + + def launch_tensor_parallel_group( + self, + server_args: ServerArgs, + port_args: PortArgs, + base_gpu_id: int, + dp_rank: int, + ): + if not server_args.enable_dp_attention: + logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.") + + memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=server_args.enable_memory_saver + ) + + scheduler_pipe_readers = [] + + nnodes_per_tp_group = max(server_args.nnodes // server_args.pp_size, 1) + tp_size_per_node = server_args.tp_size // nnodes_per_tp_group + tp_rank_range = range( + tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group), + tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1), + ) + + pp_size_per_node = max(server_args.pp_size // server_args.nnodes, 1) + pp_rank_range = range( + pp_size_per_node * (server_args.node_rank // nnodes_per_tp_group), + pp_size_per_node * (server_args.node_rank // nnodes_per_tp_group + 1), + ) + + for pp_rank in pp_rank_range: + for tp_rank in tp_rank_range: + rank_port_args = port_args + + if server_args.enable_dp_attention: + # dp attention has different sharding logic + _, _, dp_rank = compute_dp_attention_world_info( + server_args.enable_dp_attention, + tp_rank, + server_args.tp_size, + server_args.dp_size, + ) + # compute zmq ports for this dp rank + rank_port_args = PortArgs.init_new(server_args, dp_rank) + # Data parallelism reuses the tensor parallelism group, + # so all dp ranks should use the same nccl port. + rank_port_args.nccl_port = port_args.nccl_port + + reader, writer = mp.Pipe(duplex=False) + gpu_id = ( + server_args.base_gpu_id + + base_gpu_id + + ((pp_rank % pp_size_per_node) * tp_size_per_node) + + (tp_rank % tp_size_per_node) * server_args.gpu_id_step + ) + moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size) + proc = mp.Process( + target=run_scheduler_process, + args=( + server_args, + rank_port_args, + gpu_id, + tp_rank, + moe_ep_rank, + pp_rank, + dp_rank, + writer, + self.balance_meta, + ), + ) + with memory_saver_adapter.configure_subprocess(): + proc.start() + self.scheduler_procs.append(proc) + scheduler_pipe_readers.append(reader) + + # Wait for model to finish loading + scheduler_info = [] + for i in range(len(scheduler_pipe_readers)): + scheduler_info.append(scheduler_pipe_readers[i].recv()) + + self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"] + self.max_req_input_len = scheduler_info[0]["max_req_input_len"] + + def maybe_external_dp_rank_routing(self, req: Req): + if req.data_parallel_rank is not None: + logger.debug(f"Direct routing to DP rank {req.data_parallel_rank}") + self.workers[req.data_parallel_rank].send_pyobj(req) + return True + return False + + def round_robin_scheduler(self, req: Req): + if self.maybe_external_dp_rank_routing(req): + return + + if self.server_args.disaggregation_mode == "null": + self.workers[self.round_robin_counter].send_pyobj(req) + self.round_robin_counter = (self.round_robin_counter + 1) % len( + self.workers + ) + else: + self.workers[req.bootstrap_room % len(self.workers)].send_pyobj(req) + + def shortest_queue_scheduler(self, input_requests): + if self.maybe_external_dp_rank_routing(req): + return + raise NotImplementedError() + + def minimum_tokens_scheduler(self, req): + if self.maybe_external_dp_rank_routing(req): + return + + # This variable corresponds to the balance_id in TokenizedGenerateReqInput. + # We use it to to control the number of onfly tokens (requests dispatched to workers but not yet received). + def get_next_global_balance_id() -> int: + INT32_MAX = 2147483647 + current_id = self.global_balance_id + self.global_balance_id = (self.global_balance_id + 1) % INT32_MAX + return current_id + + req.dp_balance_id = get_next_global_balance_id() + with self.balance_meta.mutex: + # 1. local_tokens represents the tokens currently inferring on the worker, + # while onfly refers to the requests dispatched by the dispatcher but not yet received by the scheduler. + onfly_info = self.balance_meta.get_shared_onfly() + local_tokens = self.balance_meta.get_shared_local_tokens() + total_tokens = [ + local_token + sum(onfly_dict.values()) + for local_token, onfly_dict in zip(local_tokens, onfly_info) + ] + target_worker = total_tokens.index(min(total_tokens)) + onfly_info[target_worker][req.dp_balance_id] = len(req.input_ids) + # 2. write the new onfly info to the shm + self.balance_meta.set_shared_onfly_info(onfly_info) + + # logger.info(f"dp workers {local_tokens=}, {onfly_info=}, {target_worker=}") + self.workers[target_worker].send_pyobj(req) + + def event_loop(self): + while True: + while True: + try: + recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK) + except zmq.ZMQError: + break + + if isinstance( + recv_req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + ), + ): + self.dispatching(recv_req) + elif isinstance(recv_req, BlockReqInput): + for worker in self.workers: + worker.send_pyobj(recv_req) + else: + # Send other control messages to first worker of tp group + for worker in self.workers[:: self.control_message_step]: + worker.send_pyobj(recv_req) + + +def run_data_parallel_controller_process( + server_args: ServerArgs, + port_args: PortArgs, + pipe_writer, +): + kill_itself_when_parent_died() + setproctitle.setproctitle("sglang::data_parallel_controller") + faulthandler.enable() + configure_logger(server_args) + parent_process = psutil.Process().parent() + balance_meta = DPBalanceMeta(server_args.dp_size) + + try: + controller = DataParallelController( + server_args, port_args, dp_balance_meta=balance_meta + ) + pipe_writer.send( + { + "status": "ready", + "max_total_num_tokens": controller.max_total_num_tokens, + "max_req_input_len": controller.max_req_input_len, + } + ) + if server_args.node_rank == 0: + controller.event_loop() + for proc in controller.scheduler_procs: + proc.join() + logger.error( + f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}" + ) + except Exception: + traceback = get_exception_traceback() + logger.error(f"DataParallelController hit an exception: {traceback}") + parent_process.send_signal(signal.SIGQUIT) + finally: + # we need to destruct mp.Manager() in balance_meta + balance_meta.destructor() diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py new file mode 100644 index 000000000..bc58f4ee5 --- /dev/null +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -0,0 +1,303 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +"""DetokenizerManager is a process that detokenizes the token ids.""" + +import dataclasses +import logging +import os +import signal +from collections import OrderedDict +from typing import Dict, List, Union + +import psutil +import setproctitle +import zmq + +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.srt.managers.io_struct import ( + BatchEmbeddingOut, + BatchMultimodalDecodeReq, + BatchMultimodalOut, + BatchStrOut, + BatchTokenIDOut, + FreezeGCReq, + MultiTokenizerRegisterReq, +) +from sglang.srt.managers.multi_tokenizer_mixin import MultiHttpWorkerDetokenizerMixin +from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.utils import ( + configure_logger, + freeze_gc, + get_zmq_socket, + kill_itself_when_parent_died, +) +from sglang.utils import ( + TypeBasedDispatcher, + find_printable_text, + get_exception_traceback, +) + +logger = logging.getLogger(__name__) + +# Maximum number of request states that detokenizer can hold. When exceeded, +# oldest request states will be evicted. Default: 65536 (1<<16). +# For more details, see: https://github.com/sgl-project/sglang/issues/2812 +# Use power of 2 values for better memory allocation. +DETOKENIZER_MAX_STATES = int(os.environ.get("SGLANG_DETOKENIZER_MAX_STATES", 1 << 16)) + + +@dataclasses.dataclass +class DecodeStatus: + """Store the status of incremental decoding.""" + + decoded_text: str + decode_ids: List[int] + surr_offset: int + read_offset: int + # Offset that's sent to tokenizer for incremental update. + sent_offset: int = 0 + + +class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): + """DetokenizerManager is a process that detokenizes the token ids.""" + + def __init__( + self, + server_args: ServerArgs, + port_args: PortArgs, + ): + # Init inter-process communication + context = zmq.Context(2) + self.recv_from_scheduler = get_zmq_socket( + context, zmq.PULL, port_args.detokenizer_ipc_name, True + ) + self.send_to_tokenizer = get_zmq_socket( + context, zmq.PUSH, port_args.tokenizer_ipc_name, False + ) + + if server_args.skip_tokenizer_init: + self.tokenizer = None + else: + self.tokenizer = get_tokenizer( + server_args.tokenizer_path, + tokenizer_mode=server_args.tokenizer_mode, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + ) + + self.decode_status = LimitedCapacityDict(capacity=DETOKENIZER_MAX_STATES) + self.is_dummy = server_args.load_format == "dummy" + + self._request_dispatcher = TypeBasedDispatcher( + [ + (BatchEmbeddingOut, self.handle_batch_embedding_out), + (BatchTokenIDOut, self.handle_batch_token_id_out), + (BatchMultimodalDecodeReq, self.handle_multimodal_decode_req), + (MultiTokenizerRegisterReq, lambda x: x), + (FreezeGCReq, self.handle_freeze_gc_req), + ] + ) + + self.is_tool_call_parser_gpt_oss = server_args.tool_call_parser == "gpt-oss" + + def event_loop(self): + """The event loop that handles requests""" + while True: + recv_obj = self.recv_from_scheduler.recv_pyobj() + output = self._request_dispatcher(recv_obj) + if output is not None: + self.send_to_tokenizer.send_pyobj(output) + + def trim_matched_stop( + self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool + ): + if no_stop_trim or not finished_reason: + return output + + matched = finished_reason.get("matched", None) + if not matched: + return output + + # TODO(lmzheng): handle the case where multiple stop strs are hit + + # Trim stop str. + if isinstance(matched, str) and isinstance(output, str): + pos = output.find(matched) + return output[:pos] if pos != -1 else output + + # Trim stop token. + if isinstance(matched, int) and isinstance(output, list): + # 200012 <|call|> is the tool call token and one of eos tokens for gpt-oss model + if output[-1] == 200012 and self.is_tool_call_parser_gpt_oss: + return output + assert len(output) > 0 + return output[:-1] + return output + + def handle_batch_embedding_out(self, recv_obj: BatchEmbeddingOut): + # If it is embedding model, no detokenization is needed. + return recv_obj + + def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOut): + bs = len(recv_obj.rids) + + # Initialize decode status + read_ids, surr_ids = [], [] + for i in range(bs): + rid = recv_obj.rids[i] + if rid not in self.decode_status: + s = DecodeStatus( + decoded_text=recv_obj.decoded_texts[i], + decode_ids=recv_obj.decode_ids[i], + surr_offset=0, + read_offset=recv_obj.read_offsets[i], + ) + self.decode_status[rid] = s + else: + s = self.decode_status[rid] + s.decode_ids.extend(recv_obj.decode_ids[i]) + + read_ids.append( + self.trim_matched_stop( + s.decode_ids[s.surr_offset :], + recv_obj.finished_reasons[i], + recv_obj.no_stop_trim[i], + ) + ) + surr_ids.append(s.decode_ids[s.surr_offset : s.read_offset]) + + # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request + surr_texts = self.tokenizer.batch_decode( + surr_ids, + skip_special_tokens=recv_obj.skip_special_tokens[0], + spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0], + ) + read_texts = self.tokenizer.batch_decode( + read_ids, + skip_special_tokens=recv_obj.skip_special_tokens[0], + spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0], + ) + + # Incremental decoding + output_strs = [] + for i in range(bs): + try: + s = self.decode_status[recv_obj.rids[i]] + except KeyError: + raise RuntimeError( + f"Decode status not found for request {recv_obj.rids[i]}. " + "It may be due to the request being evicted from the decode status due to memory pressure. " + "Please increase the maximum number of requests by setting " + "the SGLANG_DETOKENIZER_MAX_STATES environment variable to a bigger value than the default value. " + f"The current value is {DETOKENIZER_MAX_STATES}. " + "For more details, see: https://github.com/sgl-project/sglang/issues/2812" + ) + new_text = read_texts[i][len(surr_texts[i]) :] + if recv_obj.finished_reasons[i] is None: + # Streaming chunk: update the decode status + if len(new_text) > 0 and not new_text.endswith("�"): + s.decoded_text = s.decoded_text + new_text + s.surr_offset = s.read_offset + s.read_offset = len(s.decode_ids) + new_text = "" + else: + new_text = find_printable_text(new_text) + + output_str = self.trim_matched_stop( + s.decoded_text + new_text, + recv_obj.finished_reasons[i], + recv_obj.no_stop_trim[i], + ) + # Incrementally send text. + incremental_output = output_str[s.sent_offset :] + s.sent_offset = len(output_str) + output_strs.append(incremental_output) + + return BatchStrOut( + rids=recv_obj.rids, + finished_reasons=recv_obj.finished_reasons, + output_strs=output_strs, + output_ids=recv_obj.decode_ids, + prompt_tokens=recv_obj.prompt_tokens, + completion_tokens=recv_obj.completion_tokens, + cached_tokens=recv_obj.cached_tokens, + spec_verify_ct=recv_obj.spec_verify_ct, + input_token_logprobs_val=recv_obj.input_token_logprobs_val, + input_token_logprobs_idx=recv_obj.input_token_logprobs_idx, + output_token_logprobs_val=recv_obj.output_token_logprobs_val, + output_token_logprobs_idx=recv_obj.output_token_logprobs_idx, + input_top_logprobs_val=recv_obj.input_top_logprobs_val, + input_top_logprobs_idx=recv_obj.input_top_logprobs_idx, + output_top_logprobs_val=recv_obj.output_top_logprobs_val, + output_top_logprobs_idx=recv_obj.output_top_logprobs_idx, + input_token_ids_logprobs_val=recv_obj.input_token_ids_logprobs_val, + input_token_ids_logprobs_idx=recv_obj.input_token_ids_logprobs_idx, + output_token_ids_logprobs_val=recv_obj.output_token_ids_logprobs_val, + output_token_ids_logprobs_idx=recv_obj.output_token_ids_logprobs_idx, + output_hidden_states=recv_obj.output_hidden_states, + placeholder_tokens_idx=None, + placeholder_tokens_val=None, + ) + + def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq): + outputs = self.tokenizer.detokenize(recv_obj) + return BatchMultimodalOut( + rids=recv_obj.rids, + finished_reasons=recv_obj.finished_reasons, + outputs=outputs, + prompt_tokens=recv_obj.prompt_tokens, + completion_tokens=recv_obj.completion_tokens, + cached_tokens=recv_obj.cached_tokens, + placeholder_tokens_idx=None, + placeholder_tokens_val=None, + ) + + def handle_freeze_gc_req(self, recv_req: FreezeGCReq): + freeze_gc("Detokenizer Manager") + return None + + +class LimitedCapacityDict(OrderedDict): + def __init__(self, capacity: int, *args, **kwargs): + super().__init__(*args, **kwargs) + self.capacity = capacity + + def __setitem__(self, key, value): + if len(self) >= self.capacity: + # Remove the oldest element (first item in the dict) + self.popitem(last=False) + # Set the new item + super().__setitem__(key, value) + + +def run_detokenizer_process( + server_args: ServerArgs, + port_args: PortArgs, +): + kill_itself_when_parent_died() + setproctitle.setproctitle("sglang::detokenizer") + configure_logger(server_args) + parent_process = psutil.Process().parent() + + try: + manager = DetokenizerManager(server_args, port_args) + if server_args.tokenizer_worker_num > 1: + manager.multi_http_worker_event_loop() + else: + manager.event_loop() + except Exception: + manager.socket_mapping.clear_all_sockets() + traceback = get_exception_traceback() + logger.error(f"DetokenizerManager hit an exception: {traceback}") + parent_process.send_signal(signal.SIGQUIT) diff --git a/python/sglang/srt/managers/disagg_service.py b/python/sglang/srt/managers/disagg_service.py new file mode 100644 index 000000000..df0eac48b --- /dev/null +++ b/python/sglang/srt/managers/disagg_service.py @@ -0,0 +1,46 @@ +"""Start bootstrap/kv-store-related server""" + +import os +from typing import Type + +from sglang.srt.disaggregation.base import BaseKVBootstrapServer +from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + KVClassType, + TransferBackend, + get_kv_class, +) +from sglang.srt.server_args import ServerArgs + + +def start_disagg_service( + server_args: ServerArgs, +): + # Start kv boostrap server on prefill + disagg_mode = DisaggregationMode(server_args.disaggregation_mode) + transfer_backend = TransferBackend(server_args.disaggregation_transfer_backend) + + if disagg_mode == DisaggregationMode.PREFILL: + # only start bootstrap server on prefill tm + kv_bootstrap_server_class: Type[BaseKVBootstrapServer] = get_kv_class( + transfer_backend, KVClassType.BOOTSTRAP_SERVER + ) + bootstrap_server: BaseKVBootstrapServer = kv_bootstrap_server_class( + host=server_args.host, + port=server_args.disaggregation_bootstrap_port, + ) + is_create_store = ( + server_args.node_rank == 0 and transfer_backend == TransferBackend.ASCEND + ) + if is_create_store: + try: + from mf_adapter import create_config_store + + ascend_url = os.getenv("ASCEND_MF_STORE_URL") + create_config_store(ascend_url) + except Exception as e: + error_message = f"Failed create mf store, invalid ascend_url." + error_message += f" With exception {e}" + raise error_message + + return bootstrap_server diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py new file mode 100644 index 000000000..6237cd383 --- /dev/null +++ b/python/sglang/srt/managers/io_struct.py @@ -0,0 +1,1332 @@ +# Copyright 2023-2024 SGLang 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. +# ============================================================================== +""" +The definition of objects transferred between different +processes (TokenizerManager, DetokenizerManager, Scheduler). +""" + +import copy +import uuid +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from sglang.srt.lora.lora_registry import LoRARef +from sglang.srt.managers.schedule_batch import BaseFinishReason +from sglang.srt.multimodal.mm_utils import has_valid_data +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.utils import ImageData + +# Handle serialization of Image for pydantic +if TYPE_CHECKING: + from PIL.Image import Image +else: + Image = Any + + +@dataclass +class SessionParams: + id: Optional[str] = None + rid: Optional[str] = None + offset: Optional[int] = None + replace: Optional[bool] = None + drop_previous_output: Optional[bool] = None + + +# Type definitions for multimodal input data +# Individual data item types for each modality +ImageDataInputItem = Union[Image, str, ImageData, Dict] +AudioDataInputItem = Union[str, Dict] +VideoDataInputItem = Union[str, Dict] +# Union type for any multimodal data item +MultimodalDataInputItem = Union[ + ImageDataInputItem, VideoDataInputItem, AudioDataInputItem +] +# Format types supporting single items, lists, or nested lists for batch processing +MultimodalDataInputFormat = Union[ + List[List[MultimodalDataInputItem]], + List[MultimodalDataInputItem], + MultimodalDataInputItem, +] + + +@dataclass +class GenerateReqInput: + # The input prompt. It can be a single prompt or a batch of prompts. + text: Optional[Union[List[str], str]] = None + # The token ids for text; one can specify either text or input_ids + input_ids: Optional[Union[List[List[int]], List[int]]] = None + # The embeddings for input_ids; one can specify either text or input_ids or input_embeds. + input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None + # The image input. It can be an image instance, file name, URL, or base64 encoded string. + # Can be formatted as: + # - Single image for a single request + # - List of images (one per request in a batch) + # - List of lists of images (multiple images per request) + # See also python/sglang/srt/utils.py:load_image for more details. + image_data: Optional[MultimodalDataInputFormat] = None + # The video input. Like image data, it can be a file name, a url, or base64 encoded string. + video_data: Optional[MultimodalDataInputFormat] = None + # The audio input. Like image data, it can be a file name, a url, or base64 encoded string. + audio_data: Optional[MultimodalDataInputFormat] = None + # The sampling_params. See descriptions below. + sampling_params: Optional[Union[List[Dict], Dict]] = None + # The request id. + rid: Optional[Union[List[str], str]] = None + # Whether to return logprobs. + return_logprob: Optional[Union[List[bool], bool]] = None + # If return logprobs, the start location in the prompt for returning logprobs. + # By default, this value is "-1", which means it will only return logprobs for output tokens. + logprob_start_len: Optional[Union[List[int], int]] = None + # If return logprobs, the number of top logprobs to return at each position. + top_logprobs_num: Optional[Union[List[int], int]] = None + # If return logprobs, the token ids to return logprob for. + token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None + # Whether to detokenize tokens in text in the returned logprobs. + return_text_in_logprobs: bool = False + # Whether to stream output. + stream: bool = False + # Whether to log metrics for this request (e.g. health_generate calls do not log metrics) + log_metrics: bool = True + # Whether to return hidden states + return_hidden_states: Union[List[bool], bool] = False + + # The modalities of the image data [image, multi-images, video] + modalities: Optional[List[str]] = None + # Session info for continual prompting + session_params: Optional[Union[List[Dict], Dict]] = None + + # The path to the LoRA adaptors + lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None + # The uid of LoRA adaptors, should be initialized by tokenizer manager + lora_id: Optional[Union[List[Optional[str]], Optional[str]]] = None + + # Custom logit processor for advanced sampling control. Must be a serialized instance + # of `CustomLogitProcessor` in python/sglang/srt/sampling/custom_logit_processor.py + # Use the processor's `to_str()` method to generate the serialized string. + custom_logit_processor: Optional[Union[List[Optional[str]], str]] = None + + # For disaggregated inference + bootstrap_host: Optional[Union[List[str], str]] = None + bootstrap_port: Optional[Union[List[Optional[int]], int]] = None + bootstrap_room: Optional[Union[List[int], int]] = None + bootstrap_pair_key: Optional[Union[List[str], str]] = None + + # For data parallel rank routing + data_parallel_rank: Optional[int] = None + + # For background responses (OpenAI responses API) + background: bool = False + + # Conversation id used for tracking requests + conversation_id: Optional[str] = None + + # Label for the request + label: Optional[str] = None + + # Priority for the request + priority: Optional[int] = None + + # Image gen grpc migration + return_bytes: bool = False + + def contains_mm_input(self) -> bool: + return ( + has_valid_data(self.image_data) + or has_valid_data(self.video_data) + or has_valid_data(self.audio_data) + ) + + def normalize_batch_and_arguments(self): + """ + Normalize the batch size and arguments for the request. + + This method resolves various input formats and ensures all parameters + are properly formatted as either single values or batches depending on the input. + It also handles parallel sampling expansion and sets default values for + unspecified parameters. + + Raises: + ValueError: If inputs are not properly specified (e.g., none or all of + text, input_ids, input_embeds are provided) + """ + self._validate_inputs() + self._determine_batch_size() + self._handle_parallel_sampling() + + if self.is_single: + self._normalize_single_inputs() + else: + self._normalize_batch_inputs() + + def _validate_inputs(self): + """Validate that the input configuration is valid.""" + if ( + self.text is None and self.input_ids is None and self.input_embeds is None + ) or ( + self.text is not None + and self.input_ids is not None + and self.input_embeds is not None + ): + raise ValueError( + "Either text, input_ids or input_embeds should be provided." + ) + + def _determine_batch_size(self): + """Determine if this is a single example or a batch and the batch size.""" + if self.text is not None: + if isinstance(self.text, str): + self.is_single = True + self.batch_size = 1 + else: + self.is_single = False + self.batch_size = len(self.text) + self.input_embeds = None + elif self.input_ids is not None: + if len(self.input_ids) == 0: + raise ValueError("input_ids cannot be empty.") + if isinstance(self.input_ids[0], int): + self.is_single = True + self.batch_size = 1 + else: + self.is_single = False + self.batch_size = len(self.input_ids) + self.input_embeds = None + else: + if isinstance(self.input_embeds[0][0], float): + self.is_single = True + self.batch_size = 1 + else: + self.is_single = False + self.batch_size = len(self.input_embeds) + + def _handle_parallel_sampling(self): + """Handle parallel sampling parameters and adjust batch size if needed.""" + # Determine parallel sample count + if self.sampling_params is None: + self.parallel_sample_num = 1 + return + elif isinstance(self.sampling_params, dict): + self.parallel_sample_num = self.sampling_params.get("n", 1) + else: # isinstance(self.sampling_params, list): + self.parallel_sample_num = self.sampling_params[0].get("n", 1) + for sampling_params in self.sampling_params: + if self.parallel_sample_num != sampling_params.get("n", 1): + raise ValueError( + "The parallel_sample_num should be the same for all samples in sample params." + ) + + # If using parallel sampling with a single example, convert to batch + if self.parallel_sample_num > 1 and self.is_single: + self.is_single = False + if self.text is not None: + self.text = [self.text] + if self.input_ids is not None: + self.input_ids = [self.input_ids] + if self.input_embeds is not None: + self.input_embeds = [self.input_embeds] + + def _normalize_single_inputs(self): + """Normalize inputs for a single example.""" + if self.sampling_params is None: + self.sampling_params = {} + if self.rid is None: + self.rid = uuid.uuid4().hex + if self.return_logprob is None: + self.return_logprob = False + if self.logprob_start_len is None: + self.logprob_start_len = -1 + if self.top_logprobs_num is None: + self.top_logprobs_num = 0 + if not self.token_ids_logprob: # covers both None and [] + self.token_ids_logprob = None + + def _normalize_batch_inputs(self): + """Normalize inputs for a batch of examples, including parallel sampling expansion.""" + # Calculate expanded batch size + if self.parallel_sample_num == 1: + num = self.batch_size + else: + # Expand parallel_sample_num + num = self.batch_size * self.parallel_sample_num + + # Expand input based on type + self._expand_inputs(num) + self._normalize_rid(num) + self._normalize_lora_paths(num) + self._normalize_image_data(num) + self._normalize_video_data(num) + self._normalize_audio_data(num) + self._normalize_sampling_params(num) + self._normalize_logprob_params(num) + self._normalize_custom_logit_processor(num) + self._normalize_bootstrap_params(num) + + def _expand_inputs(self, num): + """Expand the main inputs (text, input_ids, input_embeds) for parallel sampling.""" + if self.text is not None: + if not isinstance(self.text, list): + raise ValueError("Text should be a list for batch processing.") + self.text = self.text * self.parallel_sample_num + elif self.input_ids is not None: + if not isinstance(self.input_ids, list) or not isinstance( + self.input_ids[0], list + ): + raise ValueError( + "input_ids should be a list of lists for batch processing." + ) + self.input_ids = self.input_ids * self.parallel_sample_num + elif self.input_embeds is not None: + if not isinstance(self.input_embeds, list): + raise ValueError("input_embeds should be a list for batch processing.") + self.input_embeds = self.input_embeds * self.parallel_sample_num + + def _normalize_lora_paths(self, num): + """Normalize LoRA paths for batch processing.""" + if self.lora_path is not None: + if isinstance(self.lora_path, str): + self.lora_path = [self.lora_path] * num + elif isinstance(self.lora_path, list): + self.lora_path = self.lora_path * self.parallel_sample_num + else: + raise ValueError("lora_path should be a list or a string.") + + def _normalize_image_data(self, num): + """Normalize image data for batch processing.""" + if self.image_data is None: + self.image_data = [None] * num + elif not isinstance(self.image_data, list): + # Single image, convert to list of single-image lists + self.image_data = [[self.image_data]] * num + self.modalities = ["image"] * num + elif isinstance(self.image_data, list): + # Handle empty list case - treat as no images + if len(self.image_data) == 0: + self.image_data = [None] * num + return + + if len(self.image_data) != self.batch_size: + raise ValueError( + "The length of image_data should be equal to the batch size." + ) + + self.modalities = [] + if len(self.image_data) > 0 and isinstance(self.image_data[0], list): + # Already a list of lists, keep as is + for i in range(len(self.image_data)): + if self.image_data[i] is None or self.image_data[i] == [None]: + self.modalities.append(None) + elif len(self.image_data[i]) == 1: + self.modalities.append("image") + elif len(self.image_data[i]) > 1: + self.modalities.append("multi-images") + else: + # Ensure len(self.modalities) == len(self.image_data) + self.modalities.append(None) + # Expand parallel_sample_num + self.image_data = self.image_data * self.parallel_sample_num + self.modalities = self.modalities * self.parallel_sample_num + else: + # List of images for a batch, wrap each in a list + wrapped_images = [[img] for img in self.image_data] + # Expand for parallel sampling + self.image_data = wrapped_images * self.parallel_sample_num + self.modalities = ["image"] * num + + def _normalize_video_data(self, num): + """Normalize video data for batch processing.""" + if self.video_data is None: + self.video_data = [None] * num + elif not isinstance(self.video_data, list): + self.video_data = [self.video_data] * num + elif isinstance(self.video_data, list): + self.video_data = self.video_data * self.parallel_sample_num + + def _normalize_audio_data(self, num): + """Normalize audio data for batch processing.""" + if self.audio_data is None: + self.audio_data = [None] * num + elif not isinstance(self.audio_data, list): + self.audio_data = [self.audio_data] * num + elif isinstance(self.audio_data, list): + self.audio_data = self.audio_data * self.parallel_sample_num + + def _normalize_sampling_params(self, num): + """Normalize sampling parameters for batch processing.""" + if self.sampling_params is None: + self.sampling_params = [{}] * num + elif isinstance(self.sampling_params, dict): + self.sampling_params = [self.sampling_params] * num + else: # Already a list + self.sampling_params = self.sampling_params * self.parallel_sample_num + + def _normalize_rid(self, num): + """Normalize request IDs for batch processing.""" + if self.rid is None: + self.rid = [uuid.uuid4().hex for _ in range(num)] + elif isinstance(self.rid, str): + new_rids = [f"{self.rid}_{i}" for i in range(num)] + self.rid = new_rids + elif isinstance(self.rid, list): + # Note: the length of rid shall be the same as the batch_size, + # as the rid would be expanded for parallel sampling in tokenizer_manager + if len(self.rid) != self.batch_size: + raise ValueError( + "The specified rids length mismatch with the batch_size for batch processing." + ) + else: + raise ValueError("The rid should be a string or a list of strings.") + + def _normalize_logprob_params(self, num): + """Normalize logprob-related parameters for batch processing.""" + + # Helper function to normalize a parameter + def normalize_param(param, default_value, param_name): + if param is None: + return [default_value] * num + elif not isinstance(param, list): + return [param] * num + else: + if self.parallel_sample_num > 1: + raise ValueError( + f"Cannot use list {param_name} with parallel_sample_num > 1" + ) + return param + + # Normalize each logprob parameter + self.return_logprob = normalize_param( + self.return_logprob, False, "return_logprob" + ) + self.logprob_start_len = normalize_param( + self.logprob_start_len, -1, "logprob_start_len" + ) + self.top_logprobs_num = normalize_param( + self.top_logprobs_num, 0, "top_logprobs_num" + ) + + # Handle token_ids_logprob specially due to its nested structure + if not self.token_ids_logprob: # covers both None and [] + self.token_ids_logprob = [None] * num + elif not isinstance(self.token_ids_logprob, list): + self.token_ids_logprob = [[self.token_ids_logprob] for _ in range(num)] + elif not isinstance(self.token_ids_logprob[0], list): + self.token_ids_logprob = [ + copy.deepcopy(self.token_ids_logprob) for _ in range(num) + ] + elif self.parallel_sample_num > 1: + raise ValueError( + "Cannot use list token_ids_logprob with parallel_sample_num > 1" + ) + + def _normalize_custom_logit_processor(self, num): + """Normalize custom logit processor for batch processing.""" + if self.custom_logit_processor is None: + self.custom_logit_processor = [None] * num + elif not isinstance(self.custom_logit_processor, list): + self.custom_logit_processor = [self.custom_logit_processor] * num + elif self.parallel_sample_num > 1: + raise ValueError( + "Cannot use list custom_logit_processor with parallel_sample_num > 1" + ) + + def _normalize_bootstrap_params(self, num): + """Normalize bootstrap parameters for batch processing.""" + # Normalize bootstrap_host + if self.bootstrap_host is None: + self.bootstrap_host = [None] * num + elif not isinstance(self.bootstrap_host, list): + self.bootstrap_host = [self.bootstrap_host] * num + elif isinstance(self.bootstrap_host, list): + self.bootstrap_host = self.bootstrap_host * self.parallel_sample_num + + # Normalize bootstrap_port + if self.bootstrap_port is None: + self.bootstrap_port = [None] * num + elif not isinstance(self.bootstrap_port, list): + self.bootstrap_port = [self.bootstrap_port] * num + elif isinstance(self.bootstrap_port, list): + self.bootstrap_port = self.bootstrap_port * self.parallel_sample_num + + # Normalize bootstrap_room + if self.bootstrap_room is None: + self.bootstrap_room = [None] * num + elif not isinstance(self.bootstrap_room, list): + self.bootstrap_room = [self.bootstrap_room + i for i in range(num)] + elif isinstance(self.bootstrap_room, list): + self.bootstrap_room = self.bootstrap_room * self.parallel_sample_num + + # Normalize bootstrap_pair_key + if self.bootstrap_pair_key is None: + self.bootstrap_pair_key = [None] * num + elif not isinstance(self.bootstrap_pair_key, list): + self.bootstrap_pair_key = [self.bootstrap_pair_key] * num + elif isinstance(self.bootstrap_pair_key, list): + self.bootstrap_pair_key = self.bootstrap_pair_key * self.parallel_sample_num + + def _validate_session_params(self): + """Validate that session parameters are properly formatted.""" + if self.session_params is not None: + if not isinstance(self.session_params, dict) and not isinstance( + self.session_params[0], dict + ): + raise ValueError("Session params must be a dict or a list of dicts.") + + def regenerate_rid(self): + """Generate a new request ID and return it.""" + self.rid = uuid.uuid4().hex + return self.rid + + def __getitem__(self, i): + return GenerateReqInput( + text=self.text[i] if self.text is not None else None, + input_ids=self.input_ids[i] if self.input_ids is not None else None, + input_embeds=( + self.input_embeds[i] if self.input_embeds is not None else None + ), + image_data=self.image_data[i], + video_data=self.video_data[i], + audio_data=self.audio_data[i], + sampling_params=self.sampling_params[i], + rid=self.rid[i], + return_logprob=self.return_logprob[i], + logprob_start_len=self.logprob_start_len[i], + top_logprobs_num=self.top_logprobs_num[i], + token_ids_logprob=self.token_ids_logprob[i], + return_text_in_logprobs=self.return_text_in_logprobs, + stream=self.stream, + log_metrics=self.log_metrics, + return_hidden_states=( + self.return_hidden_states[i] + if isinstance(self.return_hidden_states, list) + else self.return_hidden_states + ), + modalities=self.modalities[i] if self.modalities else None, + session_params=self.session_params, + lora_path=self.lora_path[i] if self.lora_path is not None else None, + lora_id=self.lora_id[i] if self.lora_id is not None else None, + custom_logit_processor=( + self.custom_logit_processor[i] + if self.custom_logit_processor is not None + else None + ), + # if `__getitem__` is called, the bootstrap_host, bootstrap_port, bootstrap_room must be a list + bootstrap_host=( + self.bootstrap_host[i] if self.bootstrap_host is not None else None + ), + bootstrap_port=( + self.bootstrap_port[i] if self.bootstrap_port is not None else None + ), + bootstrap_room=( + self.bootstrap_room[i] if self.bootstrap_room is not None else None + ), + bootstrap_pair_key=( + self.bootstrap_pair_key[i] + if self.bootstrap_pair_key is not None + else None + ), + data_parallel_rank=( + self.data_parallel_rank if self.data_parallel_rank is not None else None + ), + conversation_id=self.conversation_id, + label=self.label, + priority=self.priority, + return_bytes=self.return_bytes, + ) + + +@dataclass +class TokenizedGenerateReqInput: + # The request id + rid: str + # The input text + input_text: str + # The input token ids + input_ids: List[int] + # The multimodal inputs + mm_inputs: dict + # The sampling parameters + sampling_params: SamplingParams + # Whether to return the logprobs + return_logprob: bool + # If return logprobs, the start location in the prompt for returning logprobs. + logprob_start_len: int + # If return logprobs, the number of top logprobs to return at each position. + top_logprobs_num: int + # If return logprobs, the token id to return logprob for + token_ids_logprob: List[int] + # Whether to stream output + stream: bool + # Whether to return hidden states + return_hidden_states: bool = False + + # The input embeds + input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None + + # Session info for continual prompting + session_params: Optional[SessionParams] = None + + # LoRA related + lora_id: Optional[str] = None # None means just use the base model + + # Custom logit processor for advanced sampling control. Must be a serialized instance + # of `CustomLogitProcessor` in python/sglang/srt/sampling/custom_logit_processor.py + # Use the processor's `to_str()` method to generate the serialized string. + custom_logit_processor: Optional[str] = None + + # For disaggregated inference + bootstrap_host: Optional[str] = None + bootstrap_port: Optional[int] = None + bootstrap_room: Optional[int] = None + bootstrap_pair_key: Optional[str] = None + + # For data parallel rank routing + data_parallel_rank: Optional[int] = None + + # For dp balance + dp_balance_id: int = -1 + + # Label for the request + label: Optional[str] = None + + # Priority for the request + priority: Optional[int] = None + + # Image gen grpc migration + return_bytes: bool = False + + +@dataclass +class BatchTokenizedGenerateReqInput: + # The batch of tokenized requests + batch: List[TokenizedGenerateReqInput] + + def __len__(self): + return len(self.batch) + + def __getitem__(self, i): + return self.batch[i] + + def __iter__(self): + return iter(self.batch) + + +@dataclass +class EmbeddingReqInput: + # The input prompt. It can be a single prompt or a batch of prompts. + text: Optional[Union[List[List[str]], List[str], str]] = None + # The image input. It can be an image instance, file name, URL, or base64 encoded string. + # Can be formatted as: + # - Single image for a single request + # - List of images (one per request in a batch) + # - List of lists of images (multiple images per request) + # See also python/sglang/srt/utils.py:load_image for more details. + image_data: Optional[MultimodalDataInputFormat] = None + # The video input. Like image data, it can be a file name, a url, or base64 encoded string. + video_data: Optional[MultimodalDataInputFormat] = None + # The audio input. Like image data, it can be a file name, a url, or base64 encoded string. + audio_data: Optional[MultimodalDataInputFormat] = None + # The token ids for text; one can either specify text or input_ids. + input_ids: Optional[Union[List[List[int]], List[int]]] = None + # The request id. + rid: Optional[Union[List[str], str]] = None + # Dummy sampling params for compatibility + sampling_params: Optional[Union[List[Dict], Dict]] = None + # Dummy input embeds for compatibility + input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None + # Whether to log metrics for this request (e.g. health_generate calls do not log metrics) + log_metrics: bool = True + # The modalities of the image data [image, multi-images, video] + modalities: Optional[List[str]] = None + # For cross-encoder requests + is_cross_encoder_request: bool = False + + # For background responses (OpenAI responses API) + background: bool = False + + def normalize_batch_and_arguments(self): + # at least one of text, input_ids, or image should be provided + if self.text is None and self.input_ids is None and self.image_data is None: + raise ValueError( + "At least one of text, input_ids, or image should be provided" + ) + + # text and input_ids cannot be provided at the same time + if self.text is not None and self.input_ids is not None: + raise ValueError("text and input_ids cannot be provided at the same time") + + # Derive the batch size + self.batch_size = 0 + self.is_single = True + + # check the batch size of text + if self.text is not None: + if isinstance(self.text, list): + self.batch_size += len(self.text) + self.is_single = False + else: + self.batch_size += 1 + + # check the batch size of input_ids + if self.input_ids is not None: + if isinstance(self.input_ids[0], list): + self.batch_size += len(self.input_ids) + self.is_single = False + else: + self.batch_size += 1 + + # Fill in default arguments + if self.is_single: + if self.rid is None: + self.rid = uuid.uuid4().hex + if self.sampling_params is None: + self.sampling_params = {} + self.sampling_params["max_new_tokens"] = 0 + else: + if self.rid is None: + self.rid = [uuid.uuid4().hex for _ in range(self.batch_size)] + else: + assert isinstance(self.rid, list), "The rid should be a list." + + if self.sampling_params is None: + self.sampling_params = [{}] * self.batch_size + elif isinstance(self.sampling_params, dict): + self.sampling_params = [self.sampling_params] * self.batch_size + for i in range(self.batch_size): + self.sampling_params[i]["max_new_tokens"] = 0 + + def regenerate_rid(self): + self.rid = uuid.uuid4().hex + return self.rid + + def contains_mm_input(self) -> bool: + return ( + has_valid_data(self.image_data) + or has_valid_data(self.video_data) + or has_valid_data(self.audio_data) + ) + + def __getitem__(self, i): + if self.is_cross_encoder_request: + return EmbeddingReqInput( + text=[self.text[i]] if self.text is not None else None, + sampling_params=self.sampling_params[i], + rid=self.rid[i], + is_cross_encoder_request=True, + ) + + return EmbeddingReqInput( + text=self.text[i] if self.text is not None else None, + input_ids=self.input_ids[i] if self.input_ids is not None else None, + image_data=self.image_data[i] if self.image_data is not None else None, + audio_data=self.audio_data[i] if self.audio_data is not None else None, + video_data=self.video_data[i] if self.video_data is not None else None, + sampling_params=self.sampling_params[i], + rid=self.rid[i], + ) + + +@dataclass +class TokenizedEmbeddingReqInput: + # The request id + rid: str + # The input text + input_text: str + # The input token ids + input_ids: List[int] + # The image inputs + image_inputs: dict + # The token type ids + token_type_ids: List[int] + # Dummy sampling params for compatibility + sampling_params: SamplingParams + # For data parallel rank routing + data_parallel_rank: Optional[int] = None + # For dp balance + dp_balance_id: int = -1 + + +@dataclass +class BatchTokenizedEmbeddingReqInput: + # The batch of tokenized embedding requests + batch: List[TokenizedEmbeddingReqInput] + + def __len__(self): + return len(self.batch) + + def __getitem__(self, i): + return self.batch[i] + + def __iter__(self): + return iter(self.batch) + + +@dataclass +class BatchTokenIDOut: + # The request id + rids: List[str] + # The finish reason + finished_reasons: List[BaseFinishReason] + # For incremental decoding + decoded_texts: List[str] + decode_ids: List[int] + read_offsets: List[int] + # Only used when `--skip-tokenizer-init` is on + output_ids: Optional[List[int]] + # Detokenization configs + skip_special_tokens: List[bool] + spaces_between_special_tokens: List[bool] + no_stop_trim: List[bool] + + # Token counts + prompt_tokens: List[int] + completion_tokens: List[int] + cached_tokens: List[int] + spec_verify_ct: List[int] + + # Logprobs + input_token_logprobs_val: List[float] + input_token_logprobs_idx: List[int] + output_token_logprobs_val: List[float] + output_token_logprobs_idx: List[int] + input_top_logprobs_val: List[List] + input_top_logprobs_idx: List[List] + output_top_logprobs_val: List[List] + output_top_logprobs_idx: List[List] + input_token_ids_logprobs_val: List[List] + input_token_ids_logprobs_idx: List[List] + output_token_ids_logprobs_val: List[List] + output_token_ids_logprobs_idx: List[List] + + # Hidden states + output_hidden_states: List[List[float]] + + # The information of placeholder tokens (e.g., image token) + # idx is the index of the token in the prompt after expansion. + # val is the length of padded tokens after expansion. + placeholder_tokens_idx: List[Optional[List[int]]] + placeholder_tokens_val: List[Optional[List[int]]] + + +@dataclass +class BatchMultimodalDecodeReq: + decoded_ids: List[int] + input_token_logprobs_val: List[float] + input_token_logprobs_idx: List[int] + output_token_logprobs_val: List[float] + output_token_logprobs_idx: List[int] + read_offsets: List[int] + skip_special_tokens: List[bool] + spaces_between_special_tokens: List[bool] + image_resolutions: List[List[int]] + resize_image_resolutions: List[List[int]] + + # The request id + rids: List[str] + finished_reasons: List[BaseFinishReason] + + # Token counts + prompt_tokens: List[int] + completion_tokens: List[int] + cached_tokens: List[int] + + # Placeholder token info + placeholder_tokens_idx: List[Optional[List[int]]] + placeholder_tokens_val: List[Optional[List[int]]] + + return_bytes: bool = False + + +@dataclass +class BatchStrOut: + # The request id + rids: List[str] + # The finish reason + finished_reasons: List[dict] + # The output decoded strings + output_strs: List[str] + # The token ids + output_ids: Optional[List[int]] + + # Token counts + prompt_tokens: List[int] + completion_tokens: List[int] + cached_tokens: List[int] + spec_verify_ct: List[int] + + # Logprobs + input_token_logprobs_val: List[float] + input_token_logprobs_idx: List[int] + output_token_logprobs_val: List[float] + output_token_logprobs_idx: List[int] + input_top_logprobs_val: List[List] + input_top_logprobs_idx: List[List] + output_top_logprobs_val: List[List] + output_top_logprobs_idx: List[List] + input_token_ids_logprobs_val: List[List] + input_token_ids_logprobs_idx: List[List] + output_token_ids_logprobs_val: List[List] + output_token_ids_logprobs_idx: List[List] + + # Hidden states + output_hidden_states: List[List[float]] + + placeholder_tokens_idx: List[Optional[List[int]]] + placeholder_tokens_val: List[Optional[List[int]]] + + +@dataclass +class BatchMultimodalOut: + # The request id + rids: List[str] + # The finish reason + finished_reasons: List[dict] + decoded_ids: List[List[int]] + # The outputs + outputs: Union[List[str | bytes], List[List[Dict]]] + + # probability values for input tokens and output tokens + input_token_logprobs_val: List[List[float]] + input_token_logprobs_idx: List[List[int]] + output_token_logprobs_val: List[List[float]] + output_token_logprobs_idx: List[List[int]] + + # Token counts + prompt_tokens: List[int] + completion_tokens: List[int] + cached_tokens: List[int] + + placeholder_tokens_idx: List[Optional[List[int]]] + placeholder_tokens_val: List[Optional[List[int]]] + + return_bytes: List[bool] + + +@dataclass +class BatchEmbeddingOut: + # The request id + rids: List[str] + # The finish reason + finished_reasons: List[BaseFinishReason] + # The output embedding + embeddings: List[List[float]] + # Token counts + prompt_tokens: List[int] + cached_tokens: List[int] + # Placeholder token info + placeholder_tokens_idx: List[Optional[List[int]]] + placeholder_tokens_val: List[Optional[List[int]]] + + +@dataclass +class ClearHiCacheReqInput: + pass + + +@dataclass +class ClearHiCacheReqOutput: + success: bool + + +@dataclass +class FlushCacheReqInput: + pass + + +@dataclass +class FlushCacheReqOutput: + success: bool + + +@dataclass +class UpdateWeightFromDiskReqInput: + # The model path with the new weights + model_path: str + # The format to load the weights + load_format: Optional[str] = None + # Whether to abort all requests before updating weights + abort_all_requests: bool = False + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + # Whether to update weights asynchronously + is_async: bool = False + # Whether to empty torch cache + torch_empty_cache: bool = False + # Whether to keep the scheduler paused after weight update + keep_pause: bool = False + + +@dataclass +class UpdateWeightFromDiskReqOutput: + success: bool + message: str + # Number of paused requests during weight sync. + num_paused_requests: Optional[int] = 0 + + +@dataclass +class UpdateWeightsFromDistributedReqInput: + names: List[str] + dtypes: List[str] + shapes: List[List[int]] + # The group name + group_name: str = "weight_update_group" + # Whether to flush the cache after updating weights + flush_cache: bool = True + # Whether to abort all requests before updating weights + abort_all_requests: bool = False + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + + +@dataclass +class UpdateWeightsFromDistributedReqOutput: + success: bool + message: str + + +@dataclass +class UpdateWeightsFromTensorReqInput: + """Update model weights from tensor input. + + - Tensors are serialized for transmission + - Data is structured in JSON for easy transmission over HTTP + """ + + serialized_named_tensors: List[Union[str, bytes]] + # Optional format specification for loading + load_format: Optional[str] = None + # Whether to flush the cache after updating weights + flush_cache: bool = True + # Whether to abort all requests before updating weights + abort_all_requests: bool = False + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + + +@dataclass +class UpdateWeightsFromTensorReqOutput: + success: bool + message: str + + +@dataclass +class InitWeightsUpdateGroupReqInput: + # The master address + master_address: str + # The master port + master_port: int + # The rank offset + rank_offset: int + # The world size + world_size: int + # The group name + group_name: str = "weight_update_group" + # The backend + backend: str = "nccl" + + +@dataclass +class InitWeightsUpdateGroupReqOutput: + success: bool + message: str + + +@dataclass +class UpdateWeightVersionReqInput: + # The new weight version + new_version: str + # Whether to abort all running requests before updating + abort_all_requests: bool = True + + +@dataclass +class GetWeightsByNameReqInput: + name: str + truncate_size: int = 100 + + +@dataclass +class GetWeightsByNameReqOutput: + parameter: list + + +@dataclass +class ReleaseMemoryOccupationReqInput: + # Optional tags to identify the memory region, which is primarily used for RL + # Currently we only support `weights` and `kv_cache` + tags: Optional[List[str]] = None + + +@dataclass +class ReleaseMemoryOccupationReqOutput: + pass + + +@dataclass +class ResumeMemoryOccupationReqInput: + # Optional tags to identify the memory region, which is primarily used for RL + # Currently we only support `weights` and `kv_cache` + tags: Optional[List[str]] = None + + +@dataclass +class ResumeMemoryOccupationReqOutput: + pass + + +@dataclass +class SlowDownReqInput: + forward_sleep_time: Optional[float] + + +@dataclass +class SlowDownReqOutput: + pass + + +@dataclass +class AbortReq: + # The request id + rid: str = "" + # Whether to abort all requests + abort_all: bool = False + # The finished reason data + finished_reason: Optional[Dict[str, Any]] = None + abort_reason: Optional[str] = None + # used in MultiTokenzierManager mode + rids: Optional[Union[List[str], str]] = None + + def __post_init__(self): + self.rids = self.rid + + +@dataclass +class GetInternalStateReq: + pass + + +@dataclass +class GetInternalStateReqOutput: + internal_state: Dict[Any, Any] + + +@dataclass +class SetInternalStateReq: + server_args: Dict[str, Any] + + +@dataclass +class SetInternalStateReqOutput: + updated: bool + server_args: Dict[str, Any] + + +@dataclass +class ProfileReqInput: + # The output directory + output_dir: Optional[str] = None + # If set, it profile as many as this number of steps. + # If it is set, profiling is automatically stopped after this step, and + # the caller doesn't need to run stop_profile. + start_step: Optional[int] = None + num_steps: Optional[int] = None + activities: Optional[List[str]] = None + profile_by_stage: bool = False + with_stack: Optional[bool] = None + record_shapes: Optional[bool] = None + + +class ProfileReqType(Enum): + START_PROFILE = 1 + STOP_PROFILE = 2 + + +@dataclass +class ProfileReq: + type: ProfileReqType + output_dir: Optional[str] = None + start_step: Optional[int] = None + num_steps: Optional[int] = None + activities: Optional[List[str]] = None + profile_by_stage: bool = False + with_stack: Optional[bool] = None + record_shapes: Optional[bool] = None + profile_id: Optional[str] = None + + +@dataclass +class ProfileReqOutput: + success: bool + message: str + + +@dataclass +class FreezeGCReq: + pass + + +@dataclass +class ConfigureLoggingReq: + log_requests: Optional[bool] = None + log_requests_level: Optional[int] = None + dump_requests_folder: Optional[str] = None + dump_requests_threshold: Optional[int] = None + crash_dump_folder: Optional[str] = None + + +@dataclass +class OpenSessionReqInput: + capacity_of_str_len: int + session_id: Optional[str] = None + + +@dataclass +class CloseSessionReqInput: + session_id: str + + +@dataclass +class OpenSessionReqOutput: + session_id: Optional[str] + success: bool + + +@dataclass +class HealthCheckOutput: + pass + + +class ExpertDistributionReq(Enum): + START_RECORD = 1 + STOP_RECORD = 2 + DUMP_RECORD = 3 + + +@dataclass +class ExpertDistributionReqOutput: + pass + + +@dataclass +class Function: + description: Optional[str] = None + name: Optional[str] = None + parameters: Optional[object] = None + + +@dataclass +class Tool: + function: Function + type: Optional[str] = "function" + + +@dataclass +class ParseFunctionCallReq: + text: str # The text to parse. + tools: List[Tool] = field( + default_factory=list + ) # A list of available function tools (name, parameters, etc.). + tool_call_parser: Optional[str] = ( + None # Specify the parser type, e.g. 'llama3', 'qwen25', or 'mistral'. If not specified, tries all. + ) + + +@dataclass +class SeparateReasoningReqInput: + text: str # The text to parse. + reasoning_parser: str # Specify the parser type, e.g., "deepseek-r1". + + +@dataclass +class VertexGenerateReqInput: + instances: List[dict] + parameters: Optional[dict] = None + + +@dataclass +class RpcReqInput: + method: str + parameters: Optional[Dict] = None + + +@dataclass +class RpcReqOutput: + success: bool + message: str + + +@dataclass +class LoadLoRAAdapterReqInput: + # The name of the lora module to newly loaded. + lora_name: str + # The path of loading. + lora_path: str + # Whether to pin the LoRA adapter in memory. + pinned: bool = False + # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. + lora_id: Optional[str] = None + + def to_ref(self) -> LoRARef: + return LoRARef( + lora_id=self.lora_id, + lora_name=self.lora_name, + lora_path=self.lora_path, + pinned=self.pinned, + ) + + +@dataclass +class UnloadLoRAAdapterReqInput: + # The name of lora module to unload. + lora_name: str + # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. + lora_id: Optional[str] = None + + def to_ref(self) -> LoRARef: + return LoRARef( + lora_id=self.lora_id, + lora_name=self.lora_name, + ) + + +@dataclass +class LoRAUpdateResult: + success: bool + error_message: Optional[str] = None + loaded_adapters: Optional[Dict[str, LoRARef]] = None + + +LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = LoRAUpdateResult + + +@dataclass +class MultiTokenizerRegisterReq: + rids: Optional[Union[List[str], str]] = None + ipc_name: Optional[str] = None + + +@dataclass +class MultiTokenizerWrapper: + worker_id: int + obj: Optional[Any] = None + + +class BlockReqType(Enum): + BLOCK = 1 + UNBLOCK = 2 + + +@dataclass +class BlockReqInput: + type: BlockReqType diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py new file mode 100644 index 000000000..f495904d5 --- /dev/null +++ b/python/sglang/srt/managers/mm_utils.py @@ -0,0 +1,775 @@ +""" +Multi-modality utils +""" + +import hashlib +import pickle +from abc import abstractmethod +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple + +import numpy as np +import torch +from torch import nn + +from sglang.srt.layers.multimodal import gpu_tensor_hash +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalInputs, + global_server_args_dict, +) +from sglang.srt.mem_cache.multimodal_cache import MultiModalCache +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.utils import flatten_nested_list, is_npu, print_warning_once +from sglang.utils import logger + +_is_npu = is_npu() + +# NOTE: Using the shared logger from sglang.utils instead of creating a module-specific logger +# to ensure consistent logging behavior across the codebase. This prevents issues with log +# propagation that can cause some log messages (like 'server is fired up') to not appear +# in the console when multimodal support is enabled. + +# TODO(mick): nccl +# cuda_ipc: for intranode tensor sharing +TensorTransportMode = Literal["cuda_ipc", "auto", "default"] + + +class TransportProxyTensor(torch.Tensor): + """ + A convenient torch.Tensor subclass that carries extra metadata and supports + efficient inter-process communications + """ + + @staticmethod + def __new__( + cls, + data: torch.Tensor, + name: Optional[str] = None, + fields: Optional[Dict[str, Any]] = None, + transport_mode: TensorTransportMode = "default", + *args, + **kwargs, + ): + + if not isinstance(data, torch.Tensor): + raise TypeError( + f"Input 'data' must be a torch.Tensor, but got {type(data)}" + ) + + instance = data.as_subclass(cls) + + instance._metadata = { + "name": name, + "fields": fields if fields is not None else {}, + "transport_mode": transport_mode, + } + + return instance + + def __getstate__(self): + """ + Called during pickling. Implements the serialization logic. + """ + # acquire all serialize metadata from _metadata + state = { + "metadata": self._metadata, + "tensor_data": None, + "ipc_extra": None, + } + + transport_mode = self._metadata.get("transport_mode", "default") + + if transport_mode == "cuda_ipc" and self.is_cuda: + try: + storage = self.untyped_storage() + handle = storage._share_cuda_() + + state["ipc_extra"] = { + "handle": handle, + "shape": self.shape, + "dtype": self.dtype, + "stride": self.stride(), + "device_index": self.device.index, + } + state["tensor_data"] = None + except Exception as e: + # Failed to get CUDA IPC handle (possibly tp). Falling back to default transport. + state["metadata"]["transport_mode"] = "default" + state["tensor_data"] = self.as_subclass(torch.Tensor) + else: + state["metadata"]["transport_mode"] = "default" + state["tensor_data"] = self.as_subclass(torch.Tensor) + + return state + + def __setstate__(self, state: Dict[str, Any]): + """ + Called during unpickling. Implements the deserialization logic. + """ + self._metadata = state["metadata"] + + transport_mode = self._metadata.get("transport_mode", "default") + + if transport_mode == "cuda_ipc" and state["ipc_extra"] is not None: + ipc_extra = state["ipc_extra"] + handle, shape, dtype, stride, source_device_index = ( + ipc_extra["handle"], + ipc_extra["shape"], + ipc_extra["dtype"], + ipc_extra["stride"], + ipc_extra["device_index"], + ) + + try: + target_device = torch.device(f"cuda:{source_device_index}") + with torch.cuda.device(target_device): + storage = torch.UntypedStorage._new_shared_cuda(*handle) + reconstructed_tensor = torch.empty( + 0, dtype=dtype, device=target_device + ).set_(storage, storage_offset=0, size=shape, stride=stride) + self.set_(reconstructed_tensor) + except Exception as e: + print(f"Error: Failed to deserialize from CUDA IPC handle ({e}).") + raise e + + elif state["tensor_data"] is not None: + self.set_(state["tensor_data"]) + else: + raise pickle.UnpicklingError( + "Invalid state for TransportProxyTensor: no tensor data found." + ) + + @property + def name(self) -> Optional[str]: + return self._metadata.get("name") + + @property + def fields(self) -> Dict[str, Any]: + return self._metadata.get("fields", {}) + + @property + def transport_mode(self) -> TensorTransportMode: + return self._metadata.get("transport_mode", "default") + + +class MultiModalityDataPaddingPattern: + """ + Data tokens (like image tokens) often need special handling during padding + to maintain model compatibility. This class provides the interface for + implementing different padding strategies for data tokens + """ + + @abstractmethod + def pad_input_tokens( + self, input_ids: List[int], mm_inputs: MultimodalInputs + ) -> List[int]: + """ + Pad the input ids sequence containing data tokens, and replace them with pad_values + """ + pass + + +class MultiModalityDataPaddingPatternTokenPairs(MultiModalityDataPaddingPattern): + """In this pattern, data tokens should be enclosed by special token pairs (e.g. ..., data_token_pairs) + + The padded value in a region enclosed by a token pair with be the same one, as the MultimodalDataItem's pad value + + This strategy should be applied when data content is marked by start/end token pairs in the input sequence. + """ + + def __init__( + self, + data_token_pairs: Optional[List[Tuple[int, int]]], + data_start_token_ids: Optional[List[int]] = None, + ) -> None: + """ + + Args: + data_start_token_ids marks the start of a single multimodal data + See Minicpmo's slice_start_id for example + """ + self.data_token_id_pairs = data_token_pairs + self.data_start_token_ids = data_start_token_ids or [ + s for s, _e in data_token_pairs + ] + + def pad_input_tokens( + self, input_ids: List[int], mm_inputs: MultimodalInputs + ) -> List[int]: + """ + This function will replace the data-tokens in between with pad_values accordingly + """ + pad_values = [item.pad_value for item in mm_inputs.mm_items] + data_token_pairs = self.data_token_id_pairs + mm_inputs.data_offsets = [] + if data_token_pairs is None: + data_token_pairs = [mm_inputs.im_start_id, mm_inputs.im_end_id] + if data_token_pairs is None: + print_warning_once( + "No data_token_pairs provided, RadixAttention might be influenced." + ) + return input_ids + start_token_ids = {s for s, _e in data_token_pairs} + end_tokens_ids = {e for _s, e in data_token_pairs} + + padded_ids = [] + last_idx = 0 + data_idx = -1 + + start_indices = [i for i, x in enumerate(input_ids) if x in start_token_ids] + end_indices = [i for i, x in enumerate(input_ids) if x in end_tokens_ids] + + if len(start_indices) != len(end_indices): + return input_ids + + for start_idx, end_idx in zip(start_indices, end_indices): + padded_ids.extend(input_ids[last_idx : start_idx + 1]) + + if input_ids[start_idx] in self.data_start_token_ids: + data_idx += 1 + mm_inputs.data_offsets += [start_idx] + + if data_idx >= len(pad_values): + data_idx = len(pad_values) - 1 + + num_tokens = end_idx - start_idx - 1 + pad_value = pad_values[data_idx] + padded_ids.extend([pad_value] * num_tokens) + + last_idx = end_idx + + padded_ids.extend(input_ids[last_idx:]) + + assert len(input_ids) == len(padded_ids), "Length validation fails" + return padded_ids + + +class MultiModalityDataPaddingPatternMultimodalTokens(MultiModalityDataPaddingPattern): + """In this pattern, data tokens should be represented as repetitions of a single token + e.g. ...., or